{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "010d8d52-273c-4a05-9614-b6b4a94ed9af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from typing import List\n",
    "from torch.utils.data import DataLoader\n",
    "from sentence_transformers.util import cos_sim\n",
    "from sentence_transformers import SentenceTransformer, models, losses, InputExample\n",
    "\n",
    "\n",
    "torch.set_num_threads(os.cpu_count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7cda14-f325-45c1-9188-e356d00071f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "word_embedding_model = models.Transformer('xlm-roberta-base', max_seq_length = 512)\n",
    "\n",
    "tokens = [\"[DOC]\", \"[QRY]\"]\n",
    "word_embedding_model.tokenizer.add_tokens(tokens, special_tokens=True)\n",
    "word_embedding_model.auto_model.resize_token_embeddings(len(word_embedding_model.tokenizer))\n",
    "\n",
    "pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), 'cls')\n",
    "model = SentenceTransformer(modules=[word_embedding_model, pooling_model])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcaa73cf-b51d-4aee-a655-94bc4ed9c28d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('corpus/GenQ.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "732d334a-67bb-45cb-9dde-895d0271964e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_examples = [InputExample(texts=['[QRY]'+r['query'], '[DOC]'+r['context']]) for _, r in df.iterrows()]\n",
    "train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=64)\n",
    "train_loss = losses.MultipleNegativesRankingLoss(model=model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fcf4eae-4f7a-4e66-af08-9596b9279c7e",
   "metadata": {},
   "source": [
    "train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=64)\n",
    "train_loss = losses.MultipleNegativesRankingLoss(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88b3b602-b704-44c4-9b6f-9bec6ee6df39",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    train_objectives=[(train_dataloader, train_loss)],\n",
    "    epochs=50,\n",
    "    warmup_steps=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14d30547-9009-47ca-bbfd-34736a42c32c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save('embedding')"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}