index_sparse.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load wiki"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-05T17:06:12.625841Z",
     "start_time": "2020-12-05T16:56:35.447998Z"
    },
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "from utils.data_utils import load_corpus\n",
    "\n",
    "for _ in (\"elasticsearch\", \"urllib3\"):\n",
    "    logging.getLogger(_).setLevel(logging.WARNING)\n",
    "\n",
    "hotpot_corpus, _ = load_corpus('data/hotpot-paragraph.tsv', for_hotpot=True, require_hyperlinks=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5232077\n"
     ]
    }
   ],
   "source": [
    "print(len(hotpot_corpus))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creater index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-05T17:15:40.816501Z",
     "start_time": "2020-12-05T17:15:40.538251Z"
    },
    "code_folding": [],
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'acknowledged': True,\n",
       " 'shards_acknowledged': True,\n",
       " 'index': 'enwiki-20171001-paragraph-3.1'}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from elasticsearch import Elasticsearch\n",
    "\n",
    "data_version = '3.1'\n",
    "\n",
    "es = Elasticsearch(['10.60.0.59:9200'], timeout=30)\n",
    "\n",
    "index_name = f'enwiki-20171001-paragraph-{data_version}'\n",
    "settings = {\n",
    "    \"settings\": {\n",
    "        \"analysis\": {\n",
    "            \"analyzer\": {\n",
    "                \"en_analyzer\": {  # TODO for case i'm it's\n",
    "                    \"type\": \"standard\",\n",
    "                    \"stopwords\": \"_english_\"\n",
    "                },\n",
    "                \"simple_bigram_analyzer\": {\n",
    "                    \"tokenizer\": \"standard\",\n",
    "                    \"filter\": [\"lowercase\", \"shingle\", \"asciifolding\"]\n",
    "                },\n",
    "                \"bigram_analyzer\": {\n",
    "                    \"tokenizer\": \"standard\",\n",
    "                    \"filter\": [\"lowercase\", \"stop\", \"shingle\", \"asciifolding\"]\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"doc_id\": {\n",
    "                \"type\": \"keyword\"\n",
    "            },\n",
    "            \"url\": {\n",
    "                \"type\": \"keyword\"\n",
    "            },\n",
    "            \"title\": {\n",
    "                \"type\": \"text\",\n",
    "                \"similarity\": \"BM25\",\n",
    "                \"analyzer\": \"simple\",\n",
    "                #\"copy_to\": \"full_text\",\n",
    "                \"fields\": {\n",
    "                    \"exact\": {\n",
    "                        \"type\": \"keyword\"\n",
    "                    },\n",
    "                    \"bigram\": {\n",
    "                        \"type\": \"text\",\n",
    "                        \"similarity\": \"BM25\",\n",
    "                        \"analyzer\": \"simple_bigram_analyzer\",\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"title_unescaped\": {\n",
    "                \"type\": \"text\",\n",
    "                \"similarity\": \"BM25\",\n",
    "                \"analyzer\": \"simple\",\n",
    "                #\"copy_to\": \"full_text\",\n",
    "                \"fields\": {\n",
    "                    \"exact\": {\n",
    "                        \"type\": \"keyword\"\n",
    "                    },\n",
    "                    \"bigram\": {\n",
    "                        \"type\": \"text\",\n",
    "                        \"similarity\": \"BM25\",\n",
    "                        \"analyzer\": \"simple_bigram_analyzer\",\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            \"para_id\": {\n",
    "                \"type\": \"keyword\"\n",
    "            },\n",
    "            \"para_idx\": {\n",
    "                \"type\": \"integer\"\n",
    "            },\n",
    "            \"para_num\": {\n",
    "                \"type\": \"integer\"\n",
    "            },\n",
    "            \"text\": {\n",
    "                \"type\": \"text\",\n",
    "                \"similarity\": \"BM25\",\n",
    "                \"analyzer\": \"en_analyzer\",\n",
    "                #\"copy_to\": \"full_text\",\n",
    "                \"fields\": {\n",
    "                    \"bigram\": {\n",
    "                        \"type\": \"text\",\n",
    "                        \"analyzer\": \"bigram_analyzer\"\n",
    "                    }\n",
    "                }\n",
    "            },\n",
    "            #\"full_text\":  {\n",
    "            #    \"type\": \"text\",\n",
    "            #    \"similarity\": \"BM25\",\n",
    "            #    \"analyzer\": \"en_analyzer\"\n",
    "            #},\n",
    "            \"for_hotpot\": {\n",
    "                \"type\": \"boolean\"\n",
    "            },\n",
    "            \"hyperlinks\": {\n",
    "                \"type\": \"object\",\n",
    "                \"enabled\": False\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "if es.indices.exists(index_name):\n",
    "    es.indices.delete(index_name)\n",
    "es.indices.create(index=index_name, body=settings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Index passages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-12-06T05:28:44.047352Z",
     "start_time": "2020-12-05T17:18:53.633408Z"
    },
    "code_folding": [
     14
    ]
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8777960bfadc46709fb2649304aceebd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=18231338.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total 18231338 paragraphs indexed in ES\n"
     ]
    }
   ],
   "source": [
    "from collections import defaultdict\n",
    "import json\n",
    "from html import unescape\n",
    "import time\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from elasticsearch.helpers import bulk, scan\n",
    "\n",
    "\n",
    "def act_batch(batch):\n",
    "    time_start = time.time()\n",
    "    bulk(es, batch)\n",
    "    time_end = time.time()\n",
    "    #print(f'Indexed {len(action_batch)} items, time cost: {time_end - time_start}s')\n",
    "\n",
    "\n",
    "batch_size = 1000\n",
    "\n",
    "wiki_out = open(f'data/enwiki-20171001-paragraph-{data_version}.tsv', 'w')  # include some title paragraphs for compatibility with HotpotQA\n",
    "wiki_out.write(\"id\\ttext\\ttitle\\thyperlinks\\n\")\n",
    "wiki_out.flush()\n",
    "hotpot_out = open(f'data/hotpot-paragraph-{data_version}.tsv', 'w')  # only non-empty abstract paragraphs for HotpotQA\n",
    "hotpot_out.write(\"id\\ttext\\ttitle\\thyperlinks\\tsentence_spans\\n\")\n",
    "hotpot_out.flush()\n",
    "\n",
    "total_num_para = 0\n",
    "actions = []\n",
    "# es = Elasticsearch(['10.60.0.59:9200'], timeout=30)\n",
    "query = {\"query\": {\"match_all\": {}}}\n",
    "para_num = es.count(index='enwiki-20171001-paragraph-3', body=query)['count']\n",
    "for hit in tqdm(scan(es, query=query, index='enwiki-20171001-paragraph-3', size=batch_size), total=para_num):\n",
    "    para = hit['_source']\n",
    "    new_para = {\n",
    "        \"doc_id\": para['doc_id'],\n",
    "        \"url\": para['url'],\n",
    "        \"title\": para['title'],\n",
    "        \"title_unescaped\": unescape(para['title']),\n",
    "        \"para_id\": para['para_id'],\n",
    "        \"para_idx\": para['para_idx'],\n",
    "        \"para_num\": para['para_num'],\n",
    "        \"text\": para['text'],\n",
    "        \"hyperlinks\": para['hyperlinks'],\n",
    "        \"for_hotpot\": para['for_hotpot']\n",
    "    }\n",
    "    p_id = new_para['para_id']\n",
    "    actions.append({\n",
    "        \"_index\": index_name,\n",
    "        \"_id\": p_id,\n",
    "        \"_source\": new_para\n",
    "    })\n",
    "    #if para['para_idx'] >= 0:  # not the title paragraph\n",
    "    wiki_out.write(f\"{p_id}\\t{para['text']}\\t{para['title']}\\t{json.dumps(para['hyperlinks'])}\\n\")\n",
    "    wiki_out.flush()\n",
    "    if p_id in hotpot_corpus:  # FIXME: miss some lines\n",
    "        assert para['for_hotpot']\n",
    "        hotpot_out.write(f\"{p_id}\\t{para['text']}\\t{para['title']}\\t{json.dumps(para['hyperlinks'])}\\t{hotpot_corpus[p_id]['sentence_spans']}\\n\")\n",
    "        hotpot_out.flush()\n",
    "    if len(actions) == batch_size:\n",
    "        act_batch(actions)\n",
    "        total_num_para += len(actions)\n",
    "        actions = []\n",
    "\n",
    "wiki_out.close()\n",
    "hotpot_out.close()\n",
    "if len(actions) > 0:\n",
    "    act_batch(actions)\n",
    "    total_num_para += len(actions)\n",
    "    actions = []\n",
    "\n",
    "print(f'Total {total_num_para} paragraphs indexed in ES')"
   ]
  }
 ],
 "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.6.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "217.594px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}