index_blink_elq_....ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e49d767",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlite3\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a6c5e54",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21da4f8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "con = sqlite3.connect(\"/mnt/infonas/data/harshiitb/MTP/MTP/FEVER/FEVEROUS-main/data/feverous_wikiv1.db\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da02a711",
   "metadata": {},
   "outputs": [],
   "source": [
    "cur = con.cursor()\n",
    "cur.execute(\"SELECT id FROM wiki;\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "983a94f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cur.execute(\"SELECT id FROM wiki;\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a524c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "titles = np.load(\"/mnt/infonas/data/harshiitb/MTP/MTP/data/page_ids.pkl\", allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd99686",
   "metadata": {},
   "outputs": [],
   "source": [
    "type(titles.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6db7744",
   "metadata": {},
   "outputs": [],
   "source": [
    "titles[:11]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc3be020",
   "metadata": {},
   "outputs": [],
   "source": [
    "titles = np.asarray(titles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bf8f7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "titles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee2c041d",
   "metadata": {},
   "outputs": [],
   "source": [
    "titles[2710702]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06314934",
   "metadata": {},
   "outputs": [],
   "source": [
    "data[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "805e3ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now we check BLINK here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1705394",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.cuda.set_device(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "130dcfe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import blink.main_dense as main_dense"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9f79e955",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.cuda.is_available():  True no cuda:  False\n",
      "Device found:  cuda has  1 GPUs\n",
      "torch.float32\n",
      "Inside cross_encoder and it is using:  cuda\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "\n",
    "models_path = \"/mnt/infonas/data/harshiitb/MTP/MTP/BLINK/models/\" # the path where you stored the BLINK models\n",
    "\n",
    "config = {\n",
    "    \"test_entities\": None,\n",
    "    \"test_mentions\": None,\n",
    "    \"interactive\": False,\n",
    "    \"top_k\": 10,\n",
    "    \"biencoder_model\": models_path+\"biencoder_wiki_large.bin\",\n",
    "    \"biencoder_config\": models_path+\"biencoder_wiki_large.json\",\n",
    "    \"entity_catalogue\": models_path+\"entity.jsonl\",\n",
    "    \"entity_encoding\": models_path+\"all_entities_large.t7\",\n",
    "    \"crossencoder_model\": models_path+\"crossencoder_wiki_large.bin\",\n",
    "    \"crossencoder_config\": models_path+\"crossencoder_wiki_large.json\",\n",
    "    \"fast\": False, # set this to be true if speed is a concern\n",
    "    \"output_path\": \"logs/\" # logging directory\n",
    "}\n",
    "\n",
    "args = argparse.Namespace(**config)\n",
    "\n",
    "models = main_dense.load_models(args, logger=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "473eb324",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = ['\"900\", Cahiers d\\'Italie et d\\'Europe History The first four preambles, Giustification, Basis, Advices, Analogies were published in \"French\" in the journals of autumn 1926, March and June 1927.']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b352bff",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = [\"Hi, sachin visited India\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2cf60d60",
   "metadata": {},
   "outputs": [],
   "source": [
    "import blink.ner as NER"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8ac8c81d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2022-01-15 15:54:02,840 --------------------------------------------------------------------------------\n",
      "2022-01-15 15:54:02,843 The model key 'ner-fast' now maps to 'https://huggingface.co/flair/ner-english-fast' on the HuggingFace ModelHub\n",
      "2022-01-15 15:54:02,844  - The most current version of the model is automatically downloaded from there.\n",
      "2022-01-15 15:54:02,845  - (you can alternatively manually download the original model at https://nlp.informatik.hu-berlin.de/resources/models/ner-fast/en-ner-fast-conll03-v0.4.pt)\n",
      "2022-01-15 15:54:02,846 --------------------------------------------------------------------------------\n",
      "2022-01-15 15:54:03,760 loading file /mnt/b100/d0/harshiitb/.flair/models/ner-english-fast/4c58e7191ff952c030b82db25b3694b58800b0e722ff15427f527e1631ed6142.e13c7c4664ffe2bbfa8f1f5375bd0dced866b8c1dd7ff89a6d705518abf0a611\n"
     ]
    }
   ],
   "source": [
    "ner_model = NER.get_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "69d85a3e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 736.36it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2945.44it/s]\n"
     ]
    }
   ],
   "source": [
    "annotated = main_dense._annotate(ner_model,text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "40be792a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'label': 'unknown',\n",
       "  'label_id': -1,\n",
       "  'context_left': '\"900\", ',\n",
       "  'context_right': ' the first four preambles, giustification, basis, advices, analogies were published in \"french\" in the journals of autumn 1926, march and june 1927.',\n",
       "  'mention': \"cahiers d'italie et d'europe history\",\n",
       "  'start_pos': 7,\n",
       "  'end_pos': 43,\n",
       "  'sent_idx': 0},\n",
       " {'label': 'unknown',\n",
       "  'label_id': -1,\n",
       "  'context_left': '\"900\", cahiers d\\'italie et d\\'europe history the first four preambles, ',\n",
       "  'context_right': ', basis, advices, analogies were published in \"french\" in the journals of autumn 1926, march and june 1927.',\n",
       "  'mention': 'giustification',\n",
       "  'start_pos': 70,\n",
       "  'end_pos': 84,\n",
       "  'sent_idx': 0},\n",
       " {'label': 'unknown',\n",
       "  'label_id': -1,\n",
       "  'context_left': '\"900\", cahiers d\\'italie et d\\'europe history the first four preambles, giustification, ',\n",
       "  'context_right': ', advices, analogies were published in \"french\" in the journals of autumn 1926, march and june 1927.',\n",
       "  'mention': 'basis',\n",
       "  'start_pos': 86,\n",
       "  'end_pos': 91,\n",
       "  'sent_idx': 0},\n",
       " {'label': 'unknown',\n",
       "  'label_id': -1,\n",
       "  'context_left': '\"900\", cahiers d\\'italie et d\\'europe history the first four preambles, giustification, basis, advices, ',\n",
       "  'context_right': ' were published in \"french\" in the journals of autumn 1926, march and june 1927.',\n",
       "  'mention': 'analogies',\n",
       "  'start_pos': 102,\n",
       "  'end_pos': 111,\n",
       "  'sent_idx': 0},\n",
       " {'label': 'unknown',\n",
       "  'label_id': -1,\n",
       "  'context_left': '\"900\", cahiers d\\'italie et d\\'europe history the first four preambles, giustification, basis, advices, analogies were published in \"',\n",
       "  'context_right': '\" in the journals of autumn 1926, march and june 1927.',\n",
       "  'mention': 'french',\n",
       "  'start_pos': 131,\n",
       "  'end_pos': 137,\n",
       "  'sent_idx': 0}]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "annotated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b6c2aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "window_ids = [0,1,2]\n",
    "window_density = dict.fromkeys(window_ids, 0)\n",
    "for ent in annotated:\n",
    "    window_density[ent[\"sent_idx\"]] += 1\n",
    "window_density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3423270d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  7.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Debug Step 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 1258.34it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Debug Step 2\n",
      "3/5 \r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cross Encoder about to start!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  4.79it/s]\n"
     ]
    }
   ],
   "source": [
    "_, _, _, _, _, predictions2, scores,_ = main_dense.run(args, None, *models, test_data=annotated)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "30207b4e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['History of Europe',\n",
       "  'History of France',\n",
       "  'Atlantic history',\n",
       "  'Grandes Chroniques de France',\n",
       "  'Timeline of the Middle Ages',\n",
       "  'Augustan History',\n",
       "  'History of Corsica',\n",
       "  'The Oxford History of Early Modern Europe',\n",
       "  'History of the Lombards',\n",
       "  'Fontana History of European War and Society'],\n",
       " ['Codification (law)',\n",
       "  'Codification (linguistics)',\n",
       "  'Giselle',\n",
       "  'Liberature',\n",
       "  'Glorification',\n",
       "  'Clementine literature',\n",
       "  'Gisela (name)',\n",
       "  'List of Greek and Latin roots in English/J',\n",
       "  'Giselher',\n",
       "  'Studentification'],\n",
       " ['Primary source',\n",
       "  'Historical source',\n",
       "  'Source text',\n",
       "  'Basis of estimate',\n",
       "  'Foundation (evidence)',\n",
       "  'Groundwork of the Metaphysic of Morals',\n",
       "  'Foundations of mathematics',\n",
       "  'Foundations of the Science of Knowledge',\n",
       "  'Foundations of Leninism',\n",
       "  'Basis theorem'],\n",
       " ['Analogy',\n",
       "  'Metaphor',\n",
       "  'Analogue (literature)',\n",
       "  'Imitation',\n",
       "  'Argument from analogy',\n",
       "  'Reference',\n",
       "  'Figurative analogy',\n",
       "  'Approximation',\n",
       "  'Cf.',\n",
       "  'Analogical change'],\n",
       " ['French language',\n",
       "  'France',\n",
       "  'French literature',\n",
       "  'Language',\n",
       "  'French poetry',\n",
       "  'Canadian French',\n",
       "  'Philip French',\n",
       "  '19th-century French literature',\n",
       "  'French Studies',\n",
       "  'French pronouns']]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predictions2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b65a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c0883f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "inverted_index = {}\n",
    "index = 0\n",
    "for w_id in range(len(text)):\n",
    "    for _ in range(window_density[w_id]):\n",
    "        try:\n",
    "            inverted_index[predictions[index][0]].append(w_id)\n",
    "        except KeyError:\n",
    "            inverted_index[predictions[index][0]] = []\n",
    "            inverted_index[predictions[index][0]].append(w_id)\n",
    "        index += 1\n",
    "inverted_index\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45b1806e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now try ELQ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90741922",
   "metadata": {},
   "outputs": [],
   "source": [
    "import elq.main_dense as elq_main_dense\n",
    "import sys\n",
    "import argparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "919f3d97",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path.append(\"/mnt/infonas/data/harshiitb/MTP/MTP/BLINK\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b45902",
   "metadata": {},
   "outputs": [],
   "source": [
    "sys.path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "781bafba",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_path = \"/mnt/infonas/data/harshiitb/MTP/MTP/BLINK/models/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51ee341d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4958fd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "elq_config = {\n",
    "    \"interactive\": False,\n",
    "    \"biencoder_model\": models_path+\"elq_wiki_large.bin\",\n",
    "    \"biencoder_config\": models_path+\"elq_large_params.txt\",\n",
    "    \"cand_token_ids_path\": models_path+\"entity_token_ids_128.t7\",\n",
    "    \"entity_catalogue\": models_path+\"entity.jsonl\",\n",
    "    \"entity_encoding\": models_path+\"all_entities_large.t7\",\n",
    "    \"output_path\": \"logs/\", # logging directory\n",
    "    \"faiss_index\": \"hnsw\",\n",
    "    \"index_path\": models_path+\"faiss_hnsw_index.pkl\",\n",
    "    \"num_cand_mentions\": 10,\n",
    "    \"num_cand_entities\": 10,\n",
    "    \"threshold_type\": \"joint\",\n",
    "    \"threshold\": -4.5,\n",
    "}\n",
    "\n",
    "elq_args = argparse.Namespace(**elq_config)\n",
    "\n",
    "elq_models = elq_main_dense.load_models(elq_args, logger=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38d04662",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_to_link = [{\n",
    "                    \"id\": 0,\n",
    "                    \"text\": \"paris is capital of which country?\".lower(),\n",
    "                },\n",
    "                {\n",
    "                    \"id\": 1,\n",
    "                    \"text\": \"paris is great granddaughter of whom?\".lower(),\n",
    "                },\n",
    "                {\n",
    "                    \"id\": 2,\n",
    "                    \"text\": \"Sachin Tendulkar visited India\".lower(),\n",
    "                },\n",
    "                ]\n",
    "\n",
    "predictions = elq_main_dense.run(elq_args, None, *elq_models, test_data=data_to_link)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d490d3a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6d21f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now Let's check high IDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961ad821",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5ba470f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframe = pd.read_csv(\"/mnt/infonas/data/harshiitb/MTP/MTP/data/wiki_tfidf_stems.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c8f042e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import snowballstemmer\n",
    "stemmer = snowballstemmer.stemmer('english')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a79f2938",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframe[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1406d526",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = stemmer.stemWord(\"nominated\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40ea3dca",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6533a24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = dataframe[dataframe['idf']<=7]\n",
    "df = df[df['idf']>=3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb44134",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df[[\"stem\",\"most_freq_term\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4665e74",
   "metadata": {},
   "outputs": [],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "715c702d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rare_words = df.set_index('stem').T.to_dict('list')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fc80a22",
   "metadata": {},
   "outputs": [],
   "source": [
    "rare_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f28c9689",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now load the arae_word.pkl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77455182",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/mnt/infonas/data/harshiitb/MTP/MTP/data/rare_words.pkl\", \"rb\") as file:\n",
    "    rw = pickle.load(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3702d6c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(rw.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d085d5ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "190b51f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = np.load(\"/mnt/infonas/data/harshiitb/MTP/MTP/data/idx.pkl\", allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86a2aa7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "index['Australia']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c70e8cfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "index.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a485b7e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "index1 = np.load(\"/mnt/infonas/data/harshiitb/MTP/MTP/data/idx.pkl\", allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eee23d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "index1.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6d2fa65",
   "metadata": {},
   "outputs": [],
   "source": [
    "index1['Australia']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec68c3ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "index2 = np.load(\"/mnt/infonas/data/harshiitb/MTP/MTP/data/idx1.pkl\", allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4809a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "index2.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaa50621",
   "metadata": {},
   "outputs": [],
   "source": [
    "index2['Portugal']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6cbebde",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f34ef5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flair.data import Sentence\n",
    "from flair.models import SequenceTagger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f79f9f97",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# load tagger\n",
    "tagger = SequenceTagger.load(\"flair/ner-english\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5625e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = [\"George Washington went to Washington\",\"Virat Kohli scored a century.\"]\n",
    "sens = []\n",
    "for sen in text:\n",
    "    sens.append(Sentence(sen,use_tokenizer=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c747c5ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "sens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b5b4d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "tagger.predict(sens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea960fd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sent_mentions = sens[1].to_dict(tag_type=\"ner\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "697b208b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sent_mentions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77e00d3d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}