2-1_sparse(tfidf).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TFIDF retrieval\n",
    "\n",
    "작성일자: 210116\\\n",
    "작성자: 조진욱\\\n",
    "목표: sklearn 과 scipy, nltk 를 가지고 passage retrieval 모델을 만들어보자\\\n",
    "순서: \n",
    "1. \n",
    "먼저 현재 보유한 knowledge resource 인 'test_context.json'를 tfidf 로 vectorize 해둠.\n",
    "저장안되어있다면 저장. 저장이 되어있다면 불러오기\n",
    "2. \n",
    "query 가 나열되어있는 json 형태 파일 'test_qas.json' 파일이 들어오면\n",
    "각 query 마다 가장 cosine distance 가 가까운 context 를 찾음\n",
    "3. \n",
    "그렇게 context, query pair(실제로는 정답확인을 위해 answer 까지)들을 저장해서\n",
    "retrived_test.csv 로 저장함. 이는 조금 뒤에 bert 모델이 사용하게 될 예정.\n",
    "\n",
    "비고:\n",
    "1. faiss 는 원래 sparse 에 사용하면 더 좋겠지만, 처음 retrieval 설명하는 자료라서 간단히 sorted 함수로 진행\n",
    "2. 파일 읽고 쓰는 함수는 모두 util 에 존재함\n",
    "3. retrieval 은 topk 로 바꿀 수 있음. 기본값 topk=1\n",
    "\n",
    "reference\n",
    "retrieval code는 정민수(4seaday)코드를 참고함(비공개)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.feature_extraction.text import TfidfVectorizer\n",
    "\n",
    "from scipy import spatial\n",
    "from nltk.tokenize import word_tokenize\n",
    "from nltk.stem.porter import PorterStemmer\n",
    "\n",
    "from tqdm.notebook import trange, tqdm\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import json\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import read_file, save_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[nltk_data] Downloading package punkt to /root/nltk_data...\n",
      "[nltk_data]   Unzipping tokenizers/punkt.zip.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import nltk\n",
    "nltk.download('punkt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SparseRetrieval:\n",
    "    def __init__(self, mode, data_path=None):\n",
    "        if not data_path:\n",
    "            self.data_path = cfg.squad_dir\n",
    "        else:\n",
    "            self.data_path = data_path\n",
    "        self.mode = mode\n",
    "\n",
    "    def make_embedding(self, context_name):\n",
    "        # Pickle save.\n",
    "        pickle_name = \"sparse\"+ \"_\" + self.mode + \"_embedding.bin\"\n",
    "        pickle_path = os.path.join(cfg.sparse_dir, pickle_name)\n",
    "\n",
    "        context_dict = read_file(os.path.join(self.data_path, context_name))\n",
    "        context = context_dict['text']\n",
    "\n",
    "        tfidfv = TfidfVectorizer(tokenizer=self.tokenize, ngram_range=(1,2)).fit(context)\n",
    "\n",
    "        if os.path.isfile(pickle_path):\n",
    "            with open(pickle_path, \"rb\") as file:\n",
    "                context_embeddings = pickle.load(file)\n",
    "            print(\"Embedding pickle load.\")\n",
    "        else:\n",
    "            context_embeddings = tfidfv.transform(context).toarray()\n",
    "            # Pickle save.\n",
    "            with open(pickle_path, \"wb\") as file:\n",
    "                pickle.dump(context_embeddings, file)\n",
    "            print(\"Embedding pickle saved.\")\n",
    "        return tfidfv, context_embeddings\n",
    "\n",
    "    def sparse_searching(self, sparse_query, sparse_embedding, texts, topk=1):\n",
    "        distances = spatial.distance.cdist(sparse_query, sparse_embedding, \"cosine\")[0]\n",
    "        result = zip(range(len(distances)), distances)\n",
    "        result = sorted(result, key=lambda x: x[1])\n",
    "\n",
    "        cand_dict = {}\n",
    "        candidate = []\n",
    "        cand_ids = []\n",
    "        for idx, distances in result[0:topk]: # top k\n",
    "            candidate.append(texts[idx])\n",
    "            cand_ids.append(idx)\n",
    "\n",
    "        cand_dict['text'] = candidate\n",
    "        cand_dict['ids'] = cand_ids\n",
    "        return cand_dict\n",
    "\n",
    "    def tokenize(self, text):\n",
    "        stemmer = PorterStemmer()\n",
    "        tokens = [word for word in word_tokenize(text)]\n",
    "        stems = [stemmer.stem(item) for item in tokens]\n",
    "        return stems\n",
    "\n",
    "    def retrieve(self, model, sparse_embedding, qas_filename, topk=1):\n",
    "        # 나중에 비교를 위해 context 정보도 같이 저장\n",
    "        context_dict = read_file(os.path.join(self.data_path, self.mode + \"_context.json\"))\n",
    "        context = context_dict['text']\n",
    "        context_id = context_dict['ids']\n",
    "\n",
    "        # make retrieved result as dataframe\n",
    "        que, que_id, anss, ctxs, ctx_ids = [], [], [], [], []\n",
    "\n",
    "        # load qas file.\n",
    "        qas = json.load(open(os.path.join(self.data_path, qas_filename)))['data']\n",
    "        for item in tqdm(qas):\n",
    "            query = item['question']\n",
    "            query_id = item['id']\n",
    "            answers = item['answers'][0]\n",
    "            query_s_embedding = model.transform([query]).toarray()\n",
    "            predict_dict = self.sparse_searching(query_s_embedding,\n",
    "                                                 sparse_embedding,\n",
    "                                                 context,\n",
    "                                                 topk)\n",
    "            que.append(query)\n",
    "            que_id.append(query_id)\n",
    "            anss.append(answers)\n",
    "            ctxs.append(predict_dict['text'])\n",
    "\n",
    "            tmp_ctx_ids = [context_id[predict_dict['ids'][i]] for i in range(len(predict_dict['ids']))] # retrieved 된 top k 개 context 들\n",
    "            ctx_ids.append(tmp_ctx_ids)\n",
    "\n",
    "        cqas = pd.DataFrame({\n",
    "            'question': que,\n",
    "            'que_id': que_id,\n",
    "            'answers': anss,\n",
    "            'context': ctxs,          # retrieved documents\n",
    "            'context_id': ctx_ids,\n",
    "        })\n",
    "        return cqas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import config as cfg\n",
    "mode = 'dev'\n",
    "context_file = 'dev_context.json'\n",
    "qa_file =  'dev_qa.json'\n",
    "\n",
    "ret_sp = SparseRetrieval(mode, data_path=cfg.squad_dir)\n",
    "tfidfv, context_embeddings = ret_sp.make_embedding(context_file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding pickle load.\n"
     ]
    }
   ],
   "source": [
    "# knowledge base에 있는 articles(context) 들의 정보를 임베딩해둠\n",
    "tfidfv, context_embeddings = ret_sp.make_embedding(context_file)\n",
    "cqa_df = ret_sp.retrieve(tfidfv, context_embeddings, qa_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65acd1bbc6da41dba98faed68ee53a52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(HTML(value=''), FloatProgress(value=0.0, max=204.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "res_path = os.path.join(cfg.sparse_dir, f\"retrieved_{mode}_sparse.csv\")\n",
    "cqa_df.to_csv(res_path, sep='\\t', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}