BUGFS/get_embedding.ipynb

In [1]:

from openai import OpenAI
import json
import re
import time
import tiktoken
import shelve
from tqdm.auto import tqdm

In [ ]:

client = OpenAI(
  api_key='',  # this is also the default, it can be omitted
)

In [3]:

project = "maven"

In [4]:

# functions = json.load(open("p_functions_with_times.json","r",encoding="utf-8"))
issues = json.load(open(f"{project}_issues_new.json","r"))

In [5]:

tokenizer = tiktoken.encoding_for_model("text-embedding-3-small")
def estimate_tokens_tiktoken(text):
    tokens = tokenizer.encode(text)
    return len(tokens)

def get_text_embedding(text):
    return client.embeddings.create(input = [text], model="text-embedding-3-small").data[0].embedding

def split_text(text, max_tokens=8192):
    words = text.split()
    chunks = []
    current_chunk = []
    current_tokens = 0

    for word in words:
        current_tokens += estimate_tokens_tiktoken(word + ' ')  # 단어와 공백 포함하여 토큰 계산
        if current_tokens > max_tokens:
            chunks.append(" ".join(current_chunk))
            current_chunk = []
            current_tokens = estimate_tokens_tiktoken(word + ' ')  # 새로운 청크의 첫 단어 토큰 수로 초기화
        current_chunk.append(word)

    if current_chunk:
        chunks.append(" ".join(current_chunk))

    return chunks


def get_long_text_embedding(text, max_tokens=8192, delay_per_token=0.02):
    chunks = split_text(text, max_tokens)
    # embeddings = [(len(chunk),get_text_embedding(chunk)) for chunk in chunks]
    embeddings = []
    for chunk in chunks:
        tokens = estimate_tokens_tiktoken(chunk)
        embeddings.append((len(chunk),get_text_embedding(chunk)))
        time.sleep(tokens * delay_per_token)
    return embeddings

In [6]:

# new_issues = []
# for issue in tqdm(issues):
#     # print(issue.keys())
#     if issue["title"] and issue["body"]:
#         title = re.sub(r"\s+"," ",issue["title"])
#         description = re.sub(r"\s+"," ",issue["body"])
#         embeddings = get_long_text_embedding(" ".join([title,description]))
#         # print(embeddings)
#         issue["embeddings"] = embeddings
#         new_issues.append(issue)
#     # break

# json.dump(new_issues,open(f"{project}_issues_with_text-emberdding-3-small.json","w"))

In [8]:

with shelve.open(f"{project}_issues_shelve.db") as db:
    for issue in tqdm(issues):
        # Issue 식별자 설정, 예: issue ID 사용
        issue_id = str(issue["id"])  

        # 이전에 이미 처리된 issue 건너뛰기
        if issue_id in db:
            continue

        # title과 body가 모두 있는 경우에만 처리
        if issue.get("title",None) and issue.get("body",None):
            title = re.sub(r"\s+", " ", issue["title"])
            description = re.sub(r"\s+", " ", issue["body"])
            
            # 여기서 임베딩을 계산하는 함수는 구현 필요
            embeddings = get_long_text_embedding(" ".join([title, description]))

            # Issue에 embeddings 추가 및 shelve에 저장
            issue["embeddings"] = embeddings
            db[issue_id] = issue  # shelve에 저장

# 결과를 파일로 저장
with open(f"{project}_issues_with_text-embedding-3-small.json", "w") as file:
    # shelve 데이터베이스에서 모든 저장된 issue를 읽어 JSON 파일로 저장
    with shelve.open(f"{project}_issues_shelve.db") as db:
        json.dump(list(db.values()), file)

  0%|          | 0/283 [00:00<?, ?it/s]