BUGFS/get_embedding.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import json\n",
    "import re\n",
    "import time\n",
    "import tiktoken\n",
    "import shelve\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI(\n",
    "  api_key='',  # this is also the default, it can be omitted\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "project = \"maven\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# functions = json.load(open(\"p_functions_with_times.json\",\"r\",encoding=\"utf-8\"))\n",
    "issues = json.load(open(f\"{project}_issues_new.json\",\"r\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = tiktoken.encoding_for_model(\"text-embedding-3-small\")\n",
    "def estimate_tokens_tiktoken(text):\n",
    "    tokens = tokenizer.encode(text)\n",
    "    return len(tokens)\n",
    "\n",
    "def get_text_embedding(text):\n",
    "    return client.embeddings.create(input = [text], model=\"text-embedding-3-small\").data[0].embedding\n",
    "\n",
    "def split_text(text, max_tokens=8192):\n",
    "    words = text.split()\n",
    "    chunks = []\n",
    "    current_chunk = []\n",
    "    current_tokens = 0\n",
    "\n",
    "    for word in words:\n",
    "        current_tokens += estimate_tokens_tiktoken(word + ' ')  # 단어와 공백 포함하여 토큰 계산\n",
    "        if current_tokens > max_tokens:\n",
    "            chunks.append(\" \".join(current_chunk))\n",
    "            current_chunk = []\n",
    "            current_tokens = estimate_tokens_tiktoken(word + ' ')  # 새로운 청크의 첫 단어 토큰 수로 초기화\n",
    "        current_chunk.append(word)\n",
    "\n",
    "    if current_chunk:\n",
    "        chunks.append(\" \".join(current_chunk))\n",
    "\n",
    "    return chunks\n",
    "\n",
    "\n",
    "def get_long_text_embedding(text, max_tokens=8192, delay_per_token=0.02):\n",
    "    chunks = split_text(text, max_tokens)\n",
    "    # embeddings = [(len(chunk),get_text_embedding(chunk)) for chunk in chunks]\n",
    "    embeddings = []\n",
    "    for chunk in chunks:\n",
    "        tokens = estimate_tokens_tiktoken(chunk)\n",
    "        embeddings.append((len(chunk),get_text_embedding(chunk)))\n",
    "        time.sleep(tokens * delay_per_token)\n",
    "    return embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# new_issues = []\n",
    "# for issue in tqdm(issues):\n",
    "#     # print(issue.keys())\n",
    "#     if issue[\"title\"] and issue[\"body\"]:\n",
    "#         title = re.sub(r\"\\s+\",\" \",issue[\"title\"])\n",
    "#         description = re.sub(r\"\\s+\",\" \",issue[\"body\"])\n",
    "#         embeddings = get_long_text_embedding(\" \".join([title,description]))\n",
    "#         # print(embeddings)\n",
    "#         issue[\"embeddings\"] = embeddings\n",
    "#         new_issues.append(issue)\n",
    "#     # break\n",
    "\n",
    "# json.dump(new_issues,open(f\"{project}_issues_with_text-emberdding-3-small.json\",\"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af7c03b537a9445f88a613bf355e1279",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/283 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with shelve.open(f\"{project}_issues_shelve.db\") as db:\n",
    "    for issue in tqdm(issues):\n",
    "        # Issue 식별자 설정, 예: issue ID 사용\n",
    "        issue_id = str(issue[\"id\"])  \n",
    "\n",
    "        # 이전에 이미 처리된 issue 건너뛰기\n",
    "        if issue_id in db:\n",
    "            continue\n",
    "\n",
    "        # title과 body가 모두 있는 경우에만 처리\n",
    "        if issue.get(\"title\",None) and issue.get(\"body\",None):\n",
    "            title = re.sub(r\"\\s+\", \" \", issue[\"title\"])\n",
    "            description = re.sub(r\"\\s+\", \" \", issue[\"body\"])\n",
    "            \n",
    "            # 여기서 임베딩을 계산하는 함수는 구현 필요\n",
    "            embeddings = get_long_text_embedding(\" \".join([title, description]))\n",
    "\n",
    "            # Issue에 embeddings 추가 및 shelve에 저장\n",
    "            issue[\"embeddings\"] = embeddings\n",
    "            db[issue_id] = issue  # shelve에 저장\n",
    "\n",
    "# 결과를 파일로 저장\n",
    "with open(f\"{project}_issues_with_text-embedding-3-small.json\", \"w\") as file:\n",
    "    # shelve 데이터베이스에서 모든 저장된 issue를 읽어 JSON 파일로 저장\n",
    "    with shelve.open(f\"{project}_issues_shelve.db\") as db:\n",
    "        json.dump(list(db.values()), file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}