{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07ecef6e-1c72-49a5-8cfc-4209aef67ef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import torch\n",
    "import tiktoken\n",
    "import functools\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from typing import List\n",
    "from collections import deque\n",
    "from torch.utils.data import DataLoader\n",
    "from sentence_transformers.util import cos_sim\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "\n",
    "torch.set_num_threads(os.cpu_count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d51dfdcc-009a-4f47-a3f4-922e5bb4ebba",
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = tiktoken.get_encoding(\"cl100k_base\")\n",
    "model = SentenceTransformer('embedding')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c85c3d06-362e-4d58-990f-80948886118e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(content_stream, max_tokens = 500, skip_empty = True):\n",
    "    chunk_buf = deque()\n",
    "    for line in content_stream:\n",
    "        if skip_empty and line.strip() == '':\n",
    "            continue\n",
    "        token_size = len(model.tokenizer.encode(line)) - 2 # remove start token and stop token\n",
    "\n",
    "        if sum([c[0] for c in chunk_buf]) + token_size > max_tokens and len(chunk_buf) > 0:\n",
    "            yield ''.join([c[1] for c in chunk_buf])\n",
    "            while sum([c[0] for c in chunk_buf]) + token_size > max_tokens and len(chunk_buf) > 0:\n",
    "                chunk_buf.popleft()\n",
    "        chunk_buf.append((token_size, line))\n",
    "    yield ''.join([c[1] for c in chunk_buf])\n",
    "\n",
    "def reconstruct(df):\n",
    "    title = df.title.iloc[0]\n",
    "    link = df.link.iloc[0]\n",
    "    source = df.source.iloc[0]\n",
    "    \n",
    "    chunks = encode(df.sort_values('idx').content + '\\n', 256)\n",
    "    return pd.DataFrame.from_records({\n",
    "        'title': title,\n",
    "        'link': link,\n",
    "        'content': c,\n",
    "        'source': source,\n",
    "    } for c in chunks).reset_index().rename(columns={'index': 'idx'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dd8e9b7-3012-4de6-bb31-9c138c9b5cdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.concat([\n",
    "    pd.read_csv('corpus/community.csv', keep_default_na=False),\n",
    "    pd.read_csv('corpus/kb.csv', keep_default_na=False),\n",
    "    pd.read_csv('corpus/tidb.csv', keep_default_na=False),\n",
    "    pd.read_csv('corpus/tidbcloud.csv', keep_default_na=False),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ff7b9cd-4143-486f-b2f0-e2bf9b860c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_df = df.groupby('link', group_keys=False).apply(reconstruct)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd96e7e6-5966-466f-9dda-88ee53cff8ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.cuda.device(0):\n",
    "    doc_df['vector'] = list(model.encode(('[DOC]'+doc_df.content).to_list()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c096a6a7-1f4a-4b57-aab3-130d24a262a0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "doc_df.to_pickle('corpus/documents.pkl')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}