make_conflict.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1818664a-76b0-464c-9202-950ee9136ca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset, Dataset\n",
    "import spacy\n",
    "from tqdm.auto import tqdm\n",
    "from collections import defaultdict\n",
    "from nltk import sent_tokenize\n",
    "import joblib\n",
    "#data = load_dataset(\"Seongill/nq_squad\")[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "id": "e1ea5886-6123-4331-95fe-921847d8bc74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['id', 'title', 'context', 'question', 'answers', 'masked_query', 'query_embedding'],\n",
       "    num_rows: 10\n",
       "})"
      ]
     },
     "execution_count": 274,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.select(range(10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 282,
   "id": "20c6060a-cf26-4971-9497-0324561e85c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79ef492c2d434609898d01447030372f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/98169 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "98169 98169\n"
     ]
    }
   ],
   "source": [
    "sents = []\n",
    "for row in tqdm(data):\n",
    "    context, answers = row[\"context\"], row[\"answers\"]\n",
    "    sentences = sent_tokenize(context)\n",
    "    cnt = 0\n",
    "    is_sent = False\n",
    "    for sent in sentences:\n",
    "        if is_sent:\n",
    "            break\n",
    "        for ans in answers[\"text\"]:\n",
    "            if is_sent:\n",
    "                break\n",
    "            if ans in sent:\n",
    "                is_sent = True\n",
    "                sents.append(sent)\n",
    "    if not is_sent:\n",
    "        sents.append(\"\")\n",
    "print(len(sents), len(data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 285,
   "id": "b8a7b631-9dda-45d6-ad59-c9daf0dd65f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3bfa5d3f511d48a1b6dfefbf4626fe8e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/98169 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cnt = 0\n",
    "for row, sent in zip(tqdm(data), sents):\n",
    "    answers = row[\"answers\"][\"text\"]\n",
    "    output = any([ans in sent for ans in answers])\n",
    "    if not output:\n",
    "        cnt+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 286,
   "id": "3f2e5478-d801-4188-a0ae-677a69881350",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "223\n"
     ]
    }
   ],
   "source": [
    "print(cnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 239,
   "id": "efff6cf4-5368-45d0-b7db-d286d8e2a7f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 8000\n",
    "spacy.prefer_gpu()\n",
    "nlp = spacy.load(\"en_core_web_trf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 281,
   "id": "c3183dc9-6539-426f-861d-2c84f1d26aa3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5074de81d7f14df997a461b2ca57a3f1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "docs = []\n",
    "for i in tqdm(range(0, len(sents), BATCH_SIZE)):\n",
    "    batch = sents[i:i+BATCH_SIZE]\n",
    "    docs.extend(list(nlp.pipe(batch)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 294,
   "id": "d2934773-aed5-4028-abf5-045f2ad096da",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f4bfe7a66592434ebb7ce49daec26e7d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/98169 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cnt = 0\n",
    "result, ent_type = [], []\n",
    "for doc, answers in zip(tqdm(docs),data[\"answers\"]):\n",
    "    is_ne = False\n",
    "    for ent in doc.ents:\n",
    "        if ent.text == answers[\"text\"][0]:\n",
    "            is_ne = True\n",
    "            ent_type.append(ent.label_)\n",
    "            break\n",
    "    if not is_ne:\n",
    "        ent_type.append(\"\")\n",
    "    result.append(is_ne)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 359,
   "id": "bd58445f-53a2-4e70-96da-fbe36ab8f8a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 360,
   "id": "1b6d448e-5d00-4cf5-b9e9-2a0d326cc811",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"answer_sentence\"] = sents\n",
    "df[\"is_named_entity\"] = is_ne\n",
    "df[\"ent_type\"] = ent_type\n",
    "df = df.drop_duplicates(subset=['question'], keep='first')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 361,
   "id": "dcd249b8-a5cf-4cc9-a7a4-328a1121a25d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ent_type\n",
       "CARDINAL       3982\n",
       "DATE           8127\n",
       "EVENT           479\n",
       "FAC            1101\n",
       "GPE            3716\n",
       "LANGUAGE        338\n",
       "LAW             192\n",
       "LOC             797\n",
       "MONEY           433\n",
       "NORP           1571\n",
       "ORDINAL         211\n",
       "ORG            3695\n",
       "PERCENT         984\n",
       "PERSON         6856\n",
       "PRODUCT         467\n",
       "QUANTITY        492\n",
       "TIME            107\n",
       "WORK_OF_ART    1192\n",
       "Name: id, dtype: int64"
      ]
     },
     "execution_count": 361,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = df[df[\"ent_type\"] != \"\"]\n",
    "df.groupby(\"ent_type\").count().iloc[:, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 362,
   "id": "d3fc062c-ac7c-44b1-ae55-dc4799f6fa9d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "34740"
      ]
     },
     "execution_count": 362,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb77fd8-c6f5-4743-b164-ee5b3a0b1c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "ent_to_answer = defaultdict(list)\n",
    "answers = [row[\"text\"][0] for row in df[\"answers\"].values]\n",
    "ents = df[\"ent_type\"].values.tolist()\n",
    "for ans, ent in zip(answers, ents):\n",
    "    ent_to_answer[ent].append(ans)\n",
    "ent_to_answer_unique = dict()\n",
    "for k,v in ent_to_answer.items():\n",
    "    ent_to_answer_unique[k] = list(set(v))\n",
    "cnt = 0\n",
    "for k,v in ent_to_answer_unique.items():\n",
    "    print(f'{k} {len(v)}')\n",
    "    cnt += len(v)\n",
    "print(cnt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 365,
   "id": "067972f3-e6ae-4876-88d4-6a8f14dadc71",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"num_words\"] = df[\"context\"].apply(lambda x: len(x.split()))\n",
    "df = df[df.num_words < 150]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 363,
   "id": "107aa9a6-b096-4a7f-b573-32586e2ed133",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset({\n",
      "    features: ['id', 'title', 'context', 'question', 'answers', 'masked_query', 'query_embedding', 'answer_sentence', 'is_named_entity', 'ent_type', '__index_level_0__'],\n",
      "    num_rows: 34740\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "from datasets import Dataset\n",
    "import pandas as pd\n",
    "dataset_filtered = Dataset.from_pandas(df)\n",
    "print(dataset_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 364,
   "id": "cadce548-bdfa-4bb2-8f6a-6428bec79d43",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "61477f4060974176a77862949834d920",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cfddc6742f7940fd81da6710b707c9e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/35 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "dataset_filtered.push_to_hub(\"squad_conflict_v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 378,
   "id": "c3ad4f77-368f-4425-859b-87db9e243b28",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"rewritten_context\"] = df[\"context\"].apply(lambda x: old_new_map.get(x, -1))\n",
    "df = df[df[\"rewritten_context\"] != -1]\n",
    "has_answers = []\n",
    "new_ans = []\n",
    "for r in df.iterrows():\n",
    "    row = r[1]\n",
    "    answer = row[\"answers\"][\"text\"][0]\n",
    "    new_ans.append(answer)\n",
    "    has_answer = answer in row[\"rewritten_context\"]\n",
    "    has_answers.append(has_answer)\n",
    "df[\"has_answer\"] = has_answers\n",
    "df[\"answer\"] = new_ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 379,
   "id": "35d98985-54b4-4658-9ab2-8a0939456a45",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "26559"
      ]
     },
     "execution_count": 379,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2393be91-698f-49a2-8a0c-407266633457",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_new_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "d1d0fcc2-2c0e-4fc2-a11e-678f09e74e19",
   "metadata": {},
   "outputs": [],
   "source": [
    "context_answer_map = defaultdict(list)\n",
    "for row in df.iterrows():\n",
    "    c, a = row[1][\"context\"], row[1][\"answers\"][\"text\"]\n",
    "    context_answer_map[c].extend(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "b1c0ebe4-32e2-4bb4-901a-4a13edf5c85c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "14263\n"
     ]
    }
   ],
   "source": [
    "print(len(context_answer_map.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "aa7235a8-d89b-45ca-850c-f60495894679",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['a copper statue of Christ', 'the Main Building', 'a Marian place of prayer and reflection']\n"
     ]
    }
   ],
   "source": [
    "print(context_answer_map[list(context_answer_map.keys())[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "0c0406d6-75db-435f-b953-5382a5f6a698",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Architecturally, the school has a Catholic character. Atop the Main Building\\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend \"Venite Ad Me Omnes\". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.'"
      ]
     },
     "execution_count": 89,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(context_answer_map.keys())[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "f392a867-96a6-4859-be40-cb5c1655b5b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ent_answer_map_unique = dict()\n",
    "for k, v in ent_answer_map.items():\n",
    "    ent_answer_map_unique[k] = list(set(v))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "95ca5687-255b-4a16-91c5-b4d4daea29d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ctxs, answers = list(context_answer_map.keys()), list(context_answer_map.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "286ac59a-ff06-40c5-a2d3-5909326e5eed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_prompt(context, answers):\n",
    "    return f\"Rewrite the following ariticle, maintaining its original meaning. Do not add any new information. Keep the specified words unchanged.\\nWords to Keep Unchanged: {', '.join(answers)}\\n\\nOriginal Article: {context}\\nRewritten Article:\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "28c55c23-9327-49dc-95dd-fe1ada0f3273",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14263"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompts = [make_prompt(c,a) for c, a in zip(ctxs, answers)]\n",
    "len(prompts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "f5439c3b-58b0-45cf-bdd1-c8ec17b6babd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import os\n",
    "client = OpenAI(api_key=os.environ.get(\"OPENAI_APIKEY\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "59a7d35e-5362-4b63-a58e-8ff38ca1a6a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "14263"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ctxs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "id": "85785c6b-1a8d-4573-a9bc-88537656c7d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54c8e3b6921040259a99f2d14b39c3a7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/714 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "BATCH_SIZE = 20\n",
    "preds = []\n",
    "recall = []\n",
    "for i in tqdm(range(0, len(prompts), BATCH_SIZE)):\n",
    "    batch = prompts[i:i+BATCH_SIZE]\n",
    "    batch_ans = answers[i:i+BATCH_SIZE]\n",
    "    responses = client.completions.create(\n",
    "        model=\"gpt-3.5-turbo-instruct\",\n",
    "        prompt=batch,\n",
    "        seed=42,\n",
    "        max_tokens=1024\n",
    "    )\n",
    "    for res, answer in zip(responses.choices, batch_ans):\n",
    "        pred = res.text.strip()\n",
    "        recall.append(sum([ans in pred for ans in answer])/len(answer))\n",
    "        preds.append(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "f9ef9a00-a677-4d43-b02f-180ee12287f3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['gpt_output.joblib']"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import joblib\n",
    "joblib.dump(preds, \"gpt_output.joblib\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "ac41e18c-ef69-48f9-a7c7-c47a01e5017a",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_new_map = dict()\n",
    "for ctx, pred in zip(ctxs, preds):\n",
    "    old_new_map[ctx] = pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "0f6f771b-ac8e-4959-8136-7b54f1df7bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"rewritten_context\"] = df[\"context\"].apply(lambda x: old_new_map[x])\n",
    "has_answers = []\n",
    "new_ans = []\n",
    "for r in df.iterrows():\n",
    "    row = r[1]\n",
    "    answer = row[\"answers\"][\"text\"][0]\n",
    "    new_ans.append(answer)\n",
    "    has_answer = answer in row[\"rewritten_context\"]\n",
    "    has_answers.append(has_answer)\n",
    "df[\"has_answer\"] = has_answers\n",
    "df[\"answer\"] = new_ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "37344832-f645-4e2f-87f3-ddee0aa9256e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None : 36204\n",
      "PERSON : 13847\n",
      "FAC : 949\n",
      "NORP : 1986\n",
      "DATE : 4072\n",
      "CARDINAL : 3199\n",
      "GPE : 2442\n",
      "ORG : 4424\n",
      "WORK_OF_ART : 433\n",
      "PERCENT : 750\n",
      "QUANTITY : 781\n",
      "ORDINAL : 381\n",
      "MONEY : 768\n",
      "EVENT : 462\n",
      "LANGUAGE : 132\n",
      "PRODUCT : 254\n",
      "LOC : 931\n",
      "LAW : 226\n",
      "TIME : 228\n"
     ]
    }
   ],
   "source": [
    "for k, v in ent_answer_map_unique.items():\n",
    "    print(f\"{k if k != '' else None} : {len(v)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e3949da9-f9bd-4bd8-8e26-627bbb2d4509",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "98169\n",
      "97888\n"
     ]
    }
   ],
   "source": [
    "print(len(data[\"train\"][\"question\"]))\n",
    "print(len(list(set(data[\"train\"][\"question\"]))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6756bb67-e9c9-462f-ba88-40d2256c92d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_retrieval(answers, contexts):\n",
    "    has_answers = [c[\"hasanswer\"] for c in contexts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b8061264-79d6-49c2-9cc2-d616b6c5fecf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.6829798123400626\n"
     ]
    }
   ],
   "source": [
    "recall = 0\n",
    "for row in data[\"train\"]:\n",
    "    ctxs = row[\"ctxs\"]\n",
    "    recall += int(any([c[\"hasanswer\"] for c in ctxs]))\n",
    "print(recall / len(data[\"train\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a74e0d40-3269-4ff9-93bb-fd1a90803045",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.7096952908587257\n"
     ]
    }
   ],
   "source": [
    "recall = 0\n",
    "for row in data[\"test\"]:\n",
    "    ctxs = row[\"ctxs\"]\n",
    "    recall += int(any([c[\"hasanswer\"] for c in ctxs]))\n",
    "print(recall / len(data[\"test\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e44466d-445e-4509-bdf6-34de978720b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m spacy download en_core_web_lg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 385,
   "id": "2b709e10-2cb7-4d65-ba2f-fb77c5bd5ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import spacy\n",
    "nlp = spacy.load(\"en_core_web_lg\")  # 여기서 \"en_core_web_sm\"은 선택한 모델에 따라 달라질 수 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2495c9fc-03d3-4623-80fb-2ab92be16c98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "62a88405dd654549a302f6f7942f2c58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/797 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a4af2b93798435c92c054ebd1fee6bd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3c89d87180a449c98cec8c014e7307df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/142M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb8326f4d23949e193c01108f03d7c02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c5a14430c0f440b199a8b11b56001d6d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/26799 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "20208"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#joblib.dump(ent_to_answer_unique, \"ent_to_answer_unique.joblib\")\n",
    "ent_to_answer_unique = joblib.load(\"ent_to_answer_unique.joblib\")\n",
    "data = load_dataset(\"Seongill/squad_conflict_v2_under_150\")[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6eecb463-80b5-4512-8746-e54aa8c6e5b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({'DATE': 6224,\n",
       "         'PERSON': 4933,\n",
       "         'CARDINAL': 3110,\n",
       "         'ORG': 2989,\n",
       "         'GPE': 2960,\n",
       "         'NORP': 1242,\n",
       "         'FAC': 878,\n",
       "         'WORK_OF_ART': 850,\n",
       "         'PERCENT': 806,\n",
       "         'LOC': 678,\n",
       "         'QUANTITY': 391,\n",
       "         'PRODUCT': 363,\n",
       "         'EVENT': 361,\n",
       "         'MONEY': 347,\n",
       "         'LANGUAGE': 274,\n",
       "         'ORDINAL': 174,\n",
       "         'LAW': 144,\n",
       "         'TIME': 75})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import Counter\n",
    "counter = Counter(data[\"ent_type\"])\n",
    "counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e99ac0a6-cf9a-4615-8d79-f5f7fcda12e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FAC : 933\n",
      "DATE : 3358\n",
      "CARDINAL : 1508\n",
      "GPE : 1562\n",
      "ORG : 2900\n",
      "PERSON : 4961\n",
      "WORK_OF_ART : 956\n",
      "PERCENT : 581\n",
      "QUANTITY : 463\n",
      "ORDINAL : 66\n",
      "MONEY : 368\n",
      "NORP : 835\n",
      "EVENT : 361\n",
      "LANGUAGE : 114\n",
      "PRODUCT : 414\n",
      "LOC : 563\n",
      "LAW : 171\n",
      "TIME : 94\n"
     ]
    }
   ],
   "source": [
    "for k, v in ent_to_answer_unique.items():\n",
    "    print(f\"{k} : {len(v)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "4c25b97e-e1b1-46d6-83c0-df79542ffe10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20208 19912\n"
     ]
    }
   ],
   "source": [
    "ent_to_answer_unique = joblib.load(\"ent_to_answer_unique.joblib\")\n",
    "unique_mentions = []\n",
    "mention_to_ner = dict()\n",
    "for k,v in ent_to_answer_unique.items():\n",
    "    if k == \"\":\n",
    "        print(len(v))\n",
    "        continue\n",
    "    else:\n",
    "        unique_mentions.extend(v)\n",
    "for k, v in ent_to_answer_unique.items():\n",
    "    if k == \"\":\n",
    "        continue\n",
    "    for mention in v:\n",
    "        mention_to_ner[mention] = k\n",
    "print(len(unique_mentions), len(mention_to_ner.keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "61f2d35b-c1ab-4c7a-9901-bd0601b11266",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FAC\n"
     ]
    }
   ],
   "source": [
    "for unique_mention in unique_mentions:\n",
    "    mention_to_ner[unique_mention]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "6629c3be-561d-4b7e-a887-e0898faae4b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import spacy\n",
    "nlp = spacy.load(\"en_core_web_lg\")\n",
    "spacy.prefer_gpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "43f415ae-80f7-4a8b-ae1f-d1c79f272366",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ffc2a2ee3064726bdf56f4b7ae04149",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "20208"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ent_type_vector_map = defaultdict(list)\n",
    "embeds = []\n",
    "for i in tqdm(range(0, len(unique_mentions), 2000)):\n",
    "    batch = unique_mentions[i:i+2000]\n",
    "    for mention, _doc in zip(batch, list(nlp.pipe(batch))):\n",
    "        embeds.append(_doc.vector)\n",
    "        ent_type_vector_map[mention_to_ner[mention]].append(_doc.vector)\n",
    "len(embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "75866ff1-7523-436d-91c6-37faf4a2303f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "for k, v in ent_type_vector_map.items():\n",
    "    ent_type_vector_map[k] = np.array(ent_type_vector_map[k])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ba30bf12-6d7a-4287-9345-6a260f90edbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "entity_vector_map = dict()\n",
    "for ent, vec in zip(unique_mentions, embeds):\n",
    "    entity_vector_map[ent] = vec\n",
    "# entity 별로 vector 인덱스를 만들고, 거기서 FAISS similarity search!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "1d3379f5-d240-44b6-87f8-c73bae37b39e",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "cdcde182-f7ef-43f1-802d-5256a4470d37",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.drop(\"__index_level_0__\", axis=1)\n",
    "df[\"answer\"] = df.answers.apply(lambda x: x[\"text\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "430395db-e2a2-4d38-9ee8-dbf1362321d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"random_answer\"] = df[\"answer\"].apply(lambda x: find_random_entity(mention_to_ner[x], x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "cd96e216-ac5c-4729-bc74-f4ae9185aff5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "94bcec834422402e930d68453b84affe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_918737/2546982892.py:6: RuntimeWarning: invalid value encountered in scalar divide\n",
      "  cosine_similarity = np.dot(arr1, arr2) / (norm(arr1) * norm(arr2))\n"
     ]
    }
   ],
   "source": [
    "res = []\n",
    "for row in tqdm(df.iterrows()):\n",
    "    res.append(find_non_identical_random_entity(mention_to_ner[row[1][\"answer\"]], row[1][\"answer\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "bc5cc5e7-7409-416e-a85e-ff6e138ae9ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"similar_answer\"] = res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "637ffcc6-8830-4dbc-96ae-78499d2912ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d7b17bd8a06d4bc8b4d08f32abf6a744",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/747 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90a4c4f20dd24a30a2b874541454d156",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "11451c1aa09248fa9b680afb18836222",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/230M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "21905a87ae3f4faea22b466c125f0e7e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5b35a991cb34a3bace73e97437dcb12",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/42943 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "all_df = load_dataset(\"Seongill/squad_conflict_all\")['train']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "4929fc9a-b0c9-4c64-8669-08e10d634eb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_new_map = dict()\n",
    "for row in all_df:\n",
    "    ctx, new_ctx = row[\"context\"], row[\"rewritten_context\"]\n",
    "    old_new_map[ctx] = new_ctx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "357c19b3-57ae-4fba-ab8d-e5cc7db81a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"rewritten_context\"] = df.context.apply(lambda x: old_new_map.get(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "a19795f3-76e2-4011-81ce-58d9a230b2be",
   "metadata": {},
   "outputs": [],
   "source": [
    "has_answers = []\n",
    "for r in df.iterrows():\n",
    "    row = r[1]\n",
    "    answer = row[\"answer\"]\n",
    "    has_answer = answer in row[\"rewritten_context\"]\n",
    "    has_answers.append(has_answer)\n",
    "df[\"has_answer\"] = has_answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd9c37b7-c864-4d2a-978e-b4078b508ae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[df.has_answer == True]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "670d0dab-6881-48eb-b3a3-ad263680c6dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df[df.has_answer == True]\n",
    "df = df.drop([\"is_named_entity\", \"num_words\", \"answer_sentence\"], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "d7e3cfd1-a269-4ccf-b4b2-4eb92072cbb0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>title</th>\n",
       "      <th>context</th>\n",
       "      <th>question</th>\n",
       "      <th>answers</th>\n",
       "      <th>masked_query</th>\n",
       "      <th>query_embedding</th>\n",
       "      <th>ent_type</th>\n",
       "      <th>answer</th>\n",
       "      <th>random_answer</th>\n",
       "      <th>similar_answer</th>\n",
       "      <th>rewritten_context</th>\n",
       "      <th>has_answer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5733be284776f41900661180</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>The Basilica of the Sacred heart at Notre Dame...</td>\n",
       "      <td>{'answer_start': [279], 'text': ['the Main Bui...</td>\n",
       "      <td>[MASK] at [MASK] is beside to which structure?</td>\n",
       "      <td>[0.3204971253871918, 0.20272424817085266, 0.22...</td>\n",
       "      <td>FAC</td>\n",
       "      <td>the Main Building</td>\n",
       "      <td>U.S. 1</td>\n",
       "      <td>the Royal Mews</td>\n",
       "      <td>With its Catholic character, the school's arch...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5733bed24776f41900661188</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "      <td>The university is the major seat of the Congre...</td>\n",
       "      <td>Where is the headquarters of the Congregation ...</td>\n",
       "      <td>{'answer_start': [119], 'text': ['Rome']}</td>\n",
       "      <td>Where is the headquarters of [MASK]?</td>\n",
       "      <td>[0.3285101652145386, 0.22826211154460907, 0.12...</td>\n",
       "      <td>GPE</td>\n",
       "      <td>Rome</td>\n",
       "      <td>Baden-Württemberg</td>\n",
       "      <td>Antioch</td>\n",
       "      <td>Moreau Seminary, located on the campus borderi...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5733bed24776f41900661189</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "      <td>The university is the major seat of the Congre...</td>\n",
       "      <td>What is the primary seminary of the Congregati...</td>\n",
       "      <td>{'answer_start': [145], 'text': ['Moreau Semin...</td>\n",
       "      <td>What is the primary seminary of [MASK]?</td>\n",
       "      <td>[0.1649821400642395, 0.1607779562473297, 0.091...</td>\n",
       "      <td>ORG</td>\n",
       "      <td>Moreau Seminary</td>\n",
       "      <td>The Entertainment Channel</td>\n",
       "      <td>Hofstra University</td>\n",
       "      <td>Moreau Seminary, located on the campus borderi...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5733bed24776f4190066118a</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "      <td>The university is the major seat of the Congre...</td>\n",
       "      <td>What is the oldest structure at Notre Dame?</td>\n",
       "      <td>{'answer_start': [234], 'text': ['Old College']}</td>\n",
       "      <td>What is the oldest structure at [MASK]?</td>\n",
       "      <td>[-0.02514331042766571, 0.3645743131637573, -0....</td>\n",
       "      <td>ORG</td>\n",
       "      <td>Old College</td>\n",
       "      <td>the Insurance Institute for Highway Safety</td>\n",
       "      <td>Eton College</td>\n",
       "      <td>Moreau Seminary, located on the campus borderi...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5733a6424776f41900660f51</td>\n",
       "      <td>University_of_Notre_Dame</td>\n",
       "      <td>The College of Engineering was established in ...</td>\n",
       "      <td>How many BS level degrees are offered in the C...</td>\n",
       "      <td>{'answer_start': [487], 'text': ['eight']}</td>\n",
       "      <td>How many BS level degrees are offered in [MASK...</td>\n",
       "      <td>[0.6788813471794128, 0.5213160514831543, 0.146...</td>\n",
       "      <td>CARDINAL</td>\n",
       "      <td>eight</td>\n",
       "      <td>2.1 million</td>\n",
       "      <td>over twenty</td>\n",
       "      <td>The College of Engineering was founded in 1920...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26794</th>\n",
       "      <td>57378e311c456719005744b1</td>\n",
       "      <td>Force</td>\n",
       "      <td>The origin of electric and magnetic fields wou...</td>\n",
       "      <td>How many scalar equations were formed into a s...</td>\n",
       "      <td>{'answer_start': [159, 159, 159, 159], 'text':...</td>\n",
       "      <td>How many scalar equations were formed into a s...</td>\n",
       "      <td>[0.3210879862308502, 0.09044459462165833, 0.41...</td>\n",
       "      <td>CARDINAL</td>\n",
       "      <td>20</td>\n",
       "      <td>sextuple</td>\n",
       "      <td>5 to 7</td>\n",
       "      <td>James Clerk Maxwell was responsible for fully ...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26795</th>\n",
       "      <td>57378e311c456719005744b2</td>\n",
       "      <td>Force</td>\n",
       "      <td>The origin of electric and magnetic fields wou...</td>\n",
       "      <td>How many vector equations did Heaviside and Gi...</td>\n",
       "      <td>{'answer_start': [215, 215, 215, 215], 'text':...</td>\n",
       "      <td>How many vector equations did [MASK] and Gibbs...</td>\n",
       "      <td>[-0.013793256133794785, -0.11560852080583572, ...</td>\n",
       "      <td>CARDINAL</td>\n",
       "      <td>4</td>\n",
       "      <td>approximately 119,000</td>\n",
       "      <td>5 to 7</td>\n",
       "      <td>James Clerk Maxwell was responsible for fully ...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26796</th>\n",
       "      <td>57378e311c456719005744b3</td>\n",
       "      <td>Force</td>\n",
       "      <td>The origin of electric and magnetic fields wou...</td>\n",
       "      <td>Who discovered that magnetic and electric coul...</td>\n",
       "      <td>{'answer_start': [444, 88, 444, 444], 'text': ...</td>\n",
       "      <td>Who discovered that magnetic and electric coul...</td>\n",
       "      <td>[0.1881973296403885, -0.14856640994548798, -0....</td>\n",
       "      <td>PERSON</td>\n",
       "      <td>Maxwell</td>\n",
       "      <td>Daniel Robbins</td>\n",
       "      <td>Thomas Edward Gordon</td>\n",
       "      <td>James Clerk Maxwell was responsible for fully ...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26797</th>\n",
       "      <td>57379829c3c5551400e51f41</td>\n",
       "      <td>Force</td>\n",
       "      <td>The weak force is due to the exchange of the h...</td>\n",
       "      <td>At what temperature do weak and electromagneti...</td>\n",
       "      <td>{'answer_start': [514, 501, 528, 501], 'text':...</td>\n",
       "      <td>At what temperature do weak and electromagneti...</td>\n",
       "      <td>[0.15145105123519897, -0.07466116547584534, 0....</td>\n",
       "      <td>QUANTITY</td>\n",
       "      <td>approximately 1015 kelvins</td>\n",
       "      <td>50 m</td>\n",
       "      <td>267 kilometers</td>\n",
       "      <td>The heavy W and Z bosons are responsible for t...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26798</th>\n",
       "      <td>5737a9afc3c5551400e51f64</td>\n",
       "      <td>Force</td>\n",
       "      <td>The connection between macroscopic nonconserva...</td>\n",
       "      <td>What is the law of thermodynamics associated w...</td>\n",
       "      <td>{'answer_start': [331, 331, 331, 331], 'text':...</td>\n",
       "      <td>What is the law of thermodynamics associated w...</td>\n",
       "      <td>[-0.15129241347312927, 0.08457052707672119, 0....</td>\n",
       "      <td>ORDINAL</td>\n",
       "      <td>Second</td>\n",
       "      <td>68th</td>\n",
       "      <td>second</td>\n",
       "      <td>Detailed treatment with statistical mechanics ...</td>\n",
       "      <td>True</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>25866 rows × 13 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                             id                     title  \\\n",
       "0      5733be284776f41900661180  University_of_Notre_Dame   \n",
       "1      5733bed24776f41900661188  University_of_Notre_Dame   \n",
       "2      5733bed24776f41900661189  University_of_Notre_Dame   \n",
       "3      5733bed24776f4190066118a  University_of_Notre_Dame   \n",
       "4      5733a6424776f41900660f51  University_of_Notre_Dame   \n",
       "...                         ...                       ...   \n",
       "26794  57378e311c456719005744b1                     Force   \n",
       "26795  57378e311c456719005744b2                     Force   \n",
       "26796  57378e311c456719005744b3                     Force   \n",
       "26797  57379829c3c5551400e51f41                     Force   \n",
       "26798  5737a9afc3c5551400e51f64                     Force   \n",
       "\n",
       "                                                 context  \\\n",
       "0      Architecturally, the school has a Catholic cha...   \n",
       "1      The university is the major seat of the Congre...   \n",
       "2      The university is the major seat of the Congre...   \n",
       "3      The university is the major seat of the Congre...   \n",
       "4      The College of Engineering was established in ...   \n",
       "...                                                  ...   \n",
       "26794  The origin of electric and magnetic fields wou...   \n",
       "26795  The origin of electric and magnetic fields wou...   \n",
       "26796  The origin of electric and magnetic fields wou...   \n",
       "26797  The weak force is due to the exchange of the h...   \n",
       "26798  The connection between macroscopic nonconserva...   \n",
       "\n",
       "                                                question  \\\n",
       "0      The Basilica of the Sacred heart at Notre Dame...   \n",
       "1      Where is the headquarters of the Congregation ...   \n",
       "2      What is the primary seminary of the Congregati...   \n",
       "3            What is the oldest structure at Notre Dame?   \n",
       "4      How many BS level degrees are offered in the C...   \n",
       "...                                                  ...   \n",
       "26794  How many scalar equations were formed into a s...   \n",
       "26795  How many vector equations did Heaviside and Gi...   \n",
       "26796  Who discovered that magnetic and electric coul...   \n",
       "26797  At what temperature do weak and electromagneti...   \n",
       "26798  What is the law of thermodynamics associated w...   \n",
       "\n",
       "                                                 answers  \\\n",
       "0      {'answer_start': [279], 'text': ['the Main Bui...   \n",
       "1              {'answer_start': [119], 'text': ['Rome']}   \n",
       "2      {'answer_start': [145], 'text': ['Moreau Semin...   \n",
       "3       {'answer_start': [234], 'text': ['Old College']}   \n",
       "4             {'answer_start': [487], 'text': ['eight']}   \n",
       "...                                                  ...   \n",
       "26794  {'answer_start': [159, 159, 159, 159], 'text':...   \n",
       "26795  {'answer_start': [215, 215, 215, 215], 'text':...   \n",
       "26796  {'answer_start': [444, 88, 444, 444], 'text': ...   \n",
       "26797  {'answer_start': [514, 501, 528, 501], 'text':...   \n",
       "26798  {'answer_start': [331, 331, 331, 331], 'text':...   \n",
       "\n",
       "                                            masked_query  \\\n",
       "0         [MASK] at [MASK] is beside to which structure?   \n",
       "1                   Where is the headquarters of [MASK]?   \n",
       "2                What is the primary seminary of [MASK]?   \n",
       "3                What is the oldest structure at [MASK]?   \n",
       "4      How many BS level degrees are offered in [MASK...   \n",
       "...                                                  ...   \n",
       "26794  How many scalar equations were formed into a s...   \n",
       "26795  How many vector equations did [MASK] and Gibbs...   \n",
       "26796  Who discovered that magnetic and electric coul...   \n",
       "26797  At what temperature do weak and electromagneti...   \n",
       "26798  What is the law of thermodynamics associated w...   \n",
       "\n",
       "                                         query_embedding  ent_type  \\\n",
       "0      [0.3204971253871918, 0.20272424817085266, 0.22...       FAC   \n",
       "1      [0.3285101652145386, 0.22826211154460907, 0.12...       GPE   \n",
       "2      [0.1649821400642395, 0.1607779562473297, 0.091...       ORG   \n",
       "3      [-0.02514331042766571, 0.3645743131637573, -0....       ORG   \n",
       "4      [0.6788813471794128, 0.5213160514831543, 0.146...  CARDINAL   \n",
       "...                                                  ...       ...   \n",
       "26794  [0.3210879862308502, 0.09044459462165833, 0.41...  CARDINAL   \n",
       "26795  [-0.013793256133794785, -0.11560852080583572, ...  CARDINAL   \n",
       "26796  [0.1881973296403885, -0.14856640994548798, -0....    PERSON   \n",
       "26797  [0.15145105123519897, -0.07466116547584534, 0....  QUANTITY   \n",
       "26798  [-0.15129241347312927, 0.08457052707672119, 0....   ORDINAL   \n",
       "\n",
       "                           answer                               random_answer  \\\n",
       "0               the Main Building                                      U.S. 1   \n",
       "1                            Rome                           Baden-Württemberg   \n",
       "2                 Moreau Seminary                   The Entertainment Channel   \n",
       "3                     Old College  the Insurance Institute for Highway Safety   \n",
       "4                           eight                                 2.1 million   \n",
       "...                           ...                                         ...   \n",
       "26794                          20                                    sextuple   \n",
       "26795                           4                       approximately 119,000   \n",
       "26796                     Maxwell                              Daniel Robbins   \n",
       "26797  approximately 1015 kelvins                                        50 m   \n",
       "26798                      Second                                        68th   \n",
       "\n",
       "             similar_answer  \\\n",
       "0            the Royal Mews   \n",
       "1                   Antioch   \n",
       "2        Hofstra University   \n",
       "3              Eton College   \n",
       "4               over twenty   \n",
       "...                     ...   \n",
       "26794                5 to 7   \n",
       "26795                5 to 7   \n",
       "26796  Thomas Edward Gordon   \n",
       "26797        267 kilometers   \n",
       "26798                second   \n",
       "\n",
       "                                       rewritten_context  has_answer  \n",
       "0      With its Catholic character, the school's arch...        True  \n",
       "1      Moreau Seminary, located on the campus borderi...        True  \n",
       "2      Moreau Seminary, located on the campus borderi...        True  \n",
       "3      Moreau Seminary, located on the campus borderi...        True  \n",
       "4      The College of Engineering was founded in 1920...        True  \n",
       "...                                                  ...         ...  \n",
       "26794  James Clerk Maxwell was responsible for fully ...        True  \n",
       "26795  James Clerk Maxwell was responsible for fully ...        True  \n",
       "26796  James Clerk Maxwell was responsible for fully ...        True  \n",
       "26797  The heavy W and Z bosons are responsible for t...        True  \n",
       "26798  Detailed treatment with statistical mechanics ...        True  \n",
       "\n",
       "[25866 rows x 13 columns]"
      ]
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "53945ed4-257d-4db6-91b3-4cfb1b610646",
   "metadata": {},
   "outputs": [],
   "source": [
    "ctx_to_answer = defaultdict(list)\n",
    "ctxs, answers = [], []\n",
    "for row in df[df[\"rewritten_context\"].isnull()].iterrows():\n",
    "    ctxs.append(row[1][\"context\"])\n",
    "    answers.append(row[1][\"answer\"])\n",
    "    ctx_to_answer[row[1][\"context\"]].append(row[1][\"answer\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "9fc262e7-af06-4151-ad2a-51391140fe7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_prompt(context, answers):\n",
    "    return f\"Rewrite the following ariticle, maintaining its original meaning. Do not add any new information. Keep the specified words unchanged.\\nWords to Keep Unchanged: {', '.join(answers)}\\n\\nOriginal Article: {context}\\nRewritten Article:\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "d771a04c-e061-4cc2-95da-9d4f7245eaad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "185"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompts = [make_prompt(k, v) for k, v in ctx_to_answer.items()]\n",
    "len(prompts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "5626a5fa-b6b8-48ea-845a-63d202e40172",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a2e3e2757354484a7bb3d7007b7d01a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "BATCH_SIZE = 20\n",
    "preds = []\n",
    "recall = []\n",
    "for i in tqdm(range(0, len(prompts), BATCH_SIZE)):\n",
    "    batch = prompts[i:i+BATCH_SIZE]\n",
    "    batch_ans = answers[i:i+BATCH_SIZE]\n",
    "    responses = client.completions.create(\n",
    "        model=\"gpt-3.5-turbo-instruct\",\n",
    "        prompt=batch,\n",
    "        seed=42,\n",
    "        max_tokens=1024\n",
    "    )\n",
    "    for res, answer in zip(responses.choices, batch_ans):\n",
    "        pred = res.text.strip()\n",
    "        recall.append(sum([ans in pred for ans in answer])/len(answer))\n",
    "        preds.append(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "52d89933-9729-4a52-88b8-92a50e23aa6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for pred, ctx in zip(preds, ctx_to_answer.keys()):\n",
    "    old_new_map.update({ctx:pred})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "e60b5f90-693a-4c36-8614-7483e3413c0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "final = Dataset.from_pandas(df)\n",
    "#"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "269b310d-548f-48ea-ab60-e22667f80e95",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5574b571e2f40adb7c98f11e17fce08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ddb181d8a0e64992830328938b1fa08e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/26 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "final.push_to_hub(\"squad_conflict_v2_under_150_with_substitution\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "ddf811cf-ed2a-4420-b869-6d3be3a82716",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from numpy.linalg import norm\n",
    "\n",
    "def check_cosine_similarity(arr1, arr2, thres):\n",
    "    cosine_similarity = np.dot(arr1, arr2) / (norm(arr1) * norm(arr2))\n",
    "    return (cosine_similarity <= thres) and (cosine_similarity >= 0.6)\n",
    "    \n",
    "def find_non_identical_random_entity(ent_type, origin: str, thres: float=0.8):\n",
    "    subsets = ent_to_answer_unique[ent_type]\n",
    "    #random_entities = random.sample(subsets, 100 if len(subsets)>100 else len(subsets))\n",
    "    for entity in subsets:#random_entities:\n",
    "        origin_vec, entity_vec = entity_vector_map[origin], entity_vector_map[entity]\n",
    "        if check_cosine_similarity(origin_vec, entity_vec, thres):\n",
    "            return entity\n",
    "    return random.choice(subsets)\n",
    "\n",
    "def find_random_entity(ent_type, origin: str):\n",
    "    subsets = ent_to_answer_unique[ent_type]\n",
    "    random_entity = random.choice(subsets)\n",
    "    while random_entity == origin:\n",
    "        random_entity = random.choice(subsets)\n",
    "    return random_entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 420,
   "id": "f3f17319-73ff-49b8-9f25-e6e83725661c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[QUANTITY] 800 or 1,000 ft -> 452.8 ft\n",
      "[NORP] West Flemish -> East Slavic\n",
      "[DATE] Jurassic-Cretaceous extinction -> 19th-century\n",
      "[PERSON] Zane Lowe -> RANDOM---the Abbot Suger\n",
      "[CARDINAL] 3,507 -> 70,000\n",
      "[CARDINAL] one fifth -> one hundred and fifty-seven\n",
      "[FAC] the Palazzo Pubblico -> the Theatre Royal\n",
      "[FAC] the Henry Cole Wing -> the Collegium Maius\n",
      "[PERCENT] 44.72% -> as low as 10%\n",
      "[NORP] Romanesque -> Medieval\n",
      "[PRODUCT] English Oak -> RANDOM---Zamak\n",
      "[GPE] East Prussia -> RANDOM---Sankt Goarshausen\n",
      "[DATE] December 1994 -> December 25, 2011\n",
      "[PERSON] James Green -> Dudley Simpson\n",
      "[CARDINAL] 643 million -> two million\n",
      "[PERSON] Charles IV -> RANDOM---Cannon\n",
      "[PERCENT] 2.65% -> Forty percent\n",
      "[PERSON] Kristjan Jaak Peterson -> Alvaro Martin\n",
      "[DATE] July 14, 2008 -> between 14 to 17 years of age\n",
      "[LOC] Cape Guardafui -> RANDOM---the Indian Ocean\n",
      "[ORG] NHL -> RANDOM---National Islamic Front\n",
      "[DATE] between 333 and 390 -> between 2001 and 2004\n",
      "[PERSON] Thoraya Ahmed Obaid -> Hussein Sirri Amer\n",
      "[WORK_OF_ART] Kimi ga Yo -> RANDOM---Need for Speed\n",
      "[PERSON] James Longstreet -> Dennis Bennett\n",
      "[DATE] 1995 -> June 1973\n",
      "[PERSON] Marie de' Medici -> RANDOM---James Watt\n",
      "[WORK_OF_ART] Philosophy of the Revolution -> Indiana Jones and the Last Crusade\n",
      "[PERSON] Dugald Stewart -> Godfrey Weitzel\n",
      "[DATE] 5 October 1910 -> 23 February\n",
      "[CARDINAL] 427,822 -> RANDOM---106,000\n",
      "[FAC] the Freedom Trail -> the Kodak Theatre\n",
      "[DATE] April 27, 1986 -> May 3, 2013\n",
      "[CARDINAL] Twenty-five -> 40,000-plus\n",
      "[PERCENT] 11.8% -> 65 percent\n",
      "[PERSON] Hack Wilson -> RANDOM---Athanasius the Confessor\n",
      "[PERSON] Zhou Enlai -> RANDOM---Furpaw\n",
      "[FAC] St. James' Church -> St Mary's Stadium\n",
      "[GPE] Tripoli -> Libya\n",
      "[ORG] Jim Pattison Group -> Mike Monroney Aeronautical Center\n",
      "[DATE] 1813 -> 1904\n",
      "[WORK_OF_ART] Allegro -> RANDOM---Keeping America's Promise\n",
      "[PERCENT] 86% -> 80 percent\n",
      "[ORG] Open Firmware -> RANDOM---Palmerston\n",
      "[WORK_OF_ART] The Downfall and The Death of Robert Earl of Huntington -> Six Chapters of a Floating Life\n",
      "[MONEY] $777 million -> £28 million\n",
      "[ORG] Blenheim -> RANDOM---the International Court of Justice\n",
      "[PERSON] Peter Capaldi -> Adrian Gallagher\n",
      "[PERSON] Carlo Maria di Buonaparte -> Antonio Maria Bononcini\n",
      "[FAC] JFK International Airport -> George Bush Intercontinental Airport\n"
     ]
    }
   ],
   "source": [
    "for ent in random.sample(unique_entities, 50):\n",
    "   print(f'[{mention_to_ner[ent]}] {ent} -> {find_non_identical_random_entity(mention_to_ner[ent], ent)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 427,
   "id": "ae0de6e3-e5e8-4758-af76-1604b76bff47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[EVENT] Henley Royal Regatta -> Super Bowl XLVIII\n",
      "[PERSON] Ho Keung Tse -> John XXIII\n",
      "[PERSON] Maria Skłodowska-Curie -> Arno Tausch\n",
      "[NORP] Saxon -> Mestizo\n",
      "[PERSON] Geoff Travis -> Mars\n",
      "[FAC] Avenue J -> Boulogne-sur-Mer Harbour\n",
      "[CARDINAL] 220 million -> 17 million\n",
      "[DATE] 45 BC -> 1849\n",
      "[MONEY] $25.6 billion -> $5.7 billion\n",
      "[NORP] Mohism -> Freistaaten\n",
      "[PRODUCT] Amiibo -> iMac G3\n",
      "[DATE] between 1921 and 1923 -> Between 1836 and 1842\n",
      "[PERSON] Lionel Logue -> Elbridge Gerry\n",
      "[ORG] Comics Code Authority -> Arnold, Constable\n",
      "[PERSON] Jonas Bronck -> Tom Benson\n",
      "[PERSON] Bobby Vinton -> Richard Hoggart\n",
      "[MONEY] $388 billion -> ₹1595 billion\n",
      "[PERSON] Hans Hillerbrand -> Nikolai Diletsky\n",
      "[PERSON] Anthony France -> Konstantin Chernenko\n",
      "[DATE] around 6000 BC -> 7 July\n",
      "[ORG] General Atomics -> LPGA\n",
      "[PERSON] John Paul Weier -> Amartya Sen\n",
      "[LOC] Sacramento Pass -> South-Western Ontario\n",
      "[LAW] the Defence of the Realm Act -> the Unfair Commercial Practices Directive\n",
      "[DATE] 910 -> October 11\n",
      "[PERCENT] 4.40% -> 98.1%\n",
      "[DATE] 1660 -> 2 March 1882\n",
      "[PERSON] Vishnu -> Lemming\n",
      "[ORG] El Palacio de Hierro -> The University of St. Thomas\n",
      "[GPE] Juneau -> Margraviate of Meissen\n",
      "[CARDINAL] Thirteen thousand -> 177,000\n",
      "[PERSON] Tesla -> William Morris\n",
      "[ORG] The Pardoe Theatre -> Arsenal of Democracy\n",
      "[CARDINAL] 355 -> Twenty-five\n",
      "[FAC] London Exhibition -> West 11th Street\n",
      "[PERSON] Peter I -> Flora Shaw\n",
      "[CARDINAL] over 57,000 -> 469\n",
      "[PERSON] Steven Sinofsky -> Maria Montez\n",
      "[NORP] Gramakas -> Mongols\n",
      "[DATE] November 22, 2013 -> October 2, 2008\n",
      "[ORG] Jehovah's Witnesses -> Sirindhorn International Institute of Technology\n",
      "[PERCENT] About 75% -> 14%\n",
      "[PERSON] Yuan T. Lee -> Piet Mondrian\n",
      "[NORP] Latins -> Turkic\n",
      "[ORG] DGSE -> Pinakes\n",
      "[DATE] early 2016 -> 1492\n",
      "[PERSON] József Hild -> Lawrence Lessig\n",
      "[CARDINAL] three to four -> 600\n",
      "[ORG] The Jewish Exponent -> GameTrailers\n",
      "[NORP] Galliformes -> northern Europeans\n"
     ]
    }
   ],
   "source": [
    "for ent in random.sample(unique_entities, 50):\n",
    "   print(f'[{mention_to_ner[ent]}] {ent} -> {find_random_entity(mention_to_ner[ent], ent)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "89ba4125-78eb-473e-a1ff-9cae8342c0d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read 98169 examples from /data/seongil/datasets/normalized/HF_SquadDataset_all.jsonl.gz\n"
     ]
    }
   ],
   "source": [
    "from src.classes.qadataset import QADataset\n",
    "dataset = QADataset.load(\"HF_SquadDataset_all\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e58305b8-00c9-4e0f-ab86-81148aef5fbc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset Statistics\n",
      "-------------------------------------------\n",
      "Total Examples = 98169\n",
      "Answer Type: PERSON | Size of Group: 5558\n",
      "Answer Type: None | Size of Group: 76845\n",
      "Answer Type: NUMERIC | Size of Group: 4428\n",
      "Answer Type: DATE | Size of Group: 3902\n",
      "Answer Type: LOCATION | Size of Group: 3779\n",
      "Answer Type: ORGANIZATION | Size of Group: 3657\n",
      "-------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "dataset._report_dataset_stats()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ca43a035-f1d9-4f8a-84cd-209f72cb121d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset Statistics\n",
      "-------------------------------------------\n",
      "Total Examples = 98169\n",
      "Answer Type: PERSON | Size of Group: 5558\n",
      "Answer Type: None | Size of Group: 76845\n",
      "Answer Type: NUMERIC | Size of Group: 4428\n",
      "Answer Type: DATE | Size of Group: 3902\n",
      "Answer Type: LOCATION | Size of Group: 3779\n",
      "Answer Type: ORGANIZATION | Size of Group: 3657\n",
      "-------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "grouped_examples = defaultdict(list)\n",
    "for ex in dataset.examples:\n",
    "    grouped_examples[ex.get_example_answer_type()].append(ex)\n",
    "\n",
    "print(\"Dataset Statistics\")\n",
    "print(\"-------------------------------------------\")\n",
    "print(f\"Total Examples = {len(dataset.examples)}\")\n",
    "for group, ex_list in grouped_examples.items():\n",
    "    print(f\"Answer Type: {group} | Size of Group: {len(ex_list)}\")\n",
    "print(\"-------------------------------------------\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0132f6a4-ee4f-4134-a273-8635cfb031ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "grouped_examples[\"PERSON\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55e83c55-704f-4154-ab29-759d8496d6a4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "exp",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}