full_pipeline.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Baseline\n",
        "python: 3.8.*\n",
        "\n",
        "use ```Ctrl + ]``` to collapse all section :)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download our starter pack (3~5 min)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Backup link. Use this link if the original link not working.\n",
        "# !gdown --folder 1-Gg51MsLyQQMSlrSHkgJP-sPiX1oBs0y\n",
        "\n",
        "!gdown --folder 1T6jpOtdf_i6XNYA6F_lqU4mRRh1xYPcl\n",
        "!mv baseline/* ./"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%pip install -r requirements.txt"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "h7MSEcenjVrL"
      },
      "source": [
        "notebook1\n",
        "## PART 1. Document retrieval"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Prepare the environment and import all library we need"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "niqu9pLajYC_"
      },
      "outputs": [],
      "source": [
        "# built-in libs\n",
        "import json\n",
        "import pickle\n",
        "import re\n",
        "from dataclasses import dataclass\n",
        "from pathlib import Path\n",
        "from typing import Dict, List, Set, Tuple, Union\n",
        "\n",
        "# 3rd party libs\n",
        "import hanlp\n",
        "import opencc\n",
        "import pandas as pd\n",
        "import wikipedia\n",
        "from hanlp.components.pipeline import Pipeline\n",
        "from pandarallel import pandarallel\n",
        "\n",
        "# our own libs\n",
        "from utils import load_json\n",
        "\n",
        "pandarallel.initialize(progress_bar=True, verbose=0, nb_workers=10)\n",
        "wikipedia.set_lang(\"zh\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Preload the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "TRAIN_DATA = load_json(\"data/public_train.jsonl\")\n",
        "TEST_DATA = load_json(\"data/public_test.jsonl\")\n",
        "CONVERTER_T2S = opencc.OpenCC(\"t2s.json\")\n",
        "CONVERTER_S2T = opencc.OpenCC(\"s2t.json\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Data class for type hinting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "@dataclass\n",
        "class Claim:\n",
        "    data: str\n",
        "\n",
        "@dataclass\n",
        "class AnnotationID:\n",
        "    id: int\n",
        "\n",
        "@dataclass\n",
        "class EvidenceID:\n",
        "    id: int\n",
        "\n",
        "@dataclass\n",
        "class PageTitle:\n",
        "    title: str\n",
        "\n",
        "@dataclass\n",
        "class SentenceID:\n",
        "    id: int\n",
        "\n",
        "@dataclass\n",
        "class Evidence:\n",
        "    data: List[List[Tuple[AnnotationID, EvidenceID, PageTitle, SentenceID]]]"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Helper function"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "For the sake of consistency, we convert traditional to simplified Chinese first before converting it back to traditional Chinese.  This is due to some errors occuring when converting traditional to traditional Chinese."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A3NU01DnjKp-"
      },
      "outputs": [],
      "source": [
        "def do_st_corrections(text: str) -> str:\n",
        "    simplified = CONVERTER_T2S.convert(text)\n",
        "\n",
        "    return CONVERTER_S2T.convert(simplified)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We use constituency parsing to separate part of speeches or so called constituent to extract noun phrases.  In the later stages, we will use the noun phrases as the query to search for relevant documents.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def get_nps_hanlp(\n",
        "    predictor: Pipeline,\n",
        "    d: Dict[str, Union[int, Claim, Evidence]],\n",
        ") -> List[str]:\n",
        "    claim = d[\"claim\"]\n",
        "    tree = predictor(claim)[\"con\"]\n",
        "    nps = [\n",
        "        do_st_corrections(\"\".join(subtree.leaves()))\n",
        "        for subtree in tree.subtrees(lambda t: t.label() == \"NP\")\n",
        "    ]\n",
        "\n",
        "    return nps"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Precision refers to how many related documents are retrieved.  Recall refers to how many relevant documents are retrieved.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def calculate_precision(\n",
        "    data: List[Dict[str, Union[int, Claim, Evidence]]],\n",
        "    predictions: pd.Series,\n",
        ") -> None:\n",
        "    precision = 0\n",
        "    count = 0\n",
        "\n",
        "    for i, d in enumerate(data):\n",
        "        if d[\"label\"] == \"NOT ENOUGH INFO\":\n",
        "            continue\n",
        "\n",
        "        # Extract all ground truth of titles of the wikipedia pages\n",
        "        # evidence[2] refers to the title of the wikipedia page\n",
        "        gt_pages = set([\n",
        "            evidence[2]\n",
        "            for evidence_set in d[\"evidence\"]\n",
        "            for evidence in evidence_set\n",
        "        ])\n",
        "\n",
        "        predicted_pages = predictions.iloc[i]\n",
        "        hits = predicted_pages.intersection(gt_pages)\n",
        "        if len(predicted_pages) != 0:\n",
        "            precision += len(hits) / len(predicted_pages)\n",
        "\n",
        "        count += 1\n",
        "\n",
        "    # Macro precision\n",
        "    print(f\"Precision: {precision / count}\")\n",
        "\n",
        "\n",
        "def calculate_recall(\n",
        "    data: List[Dict[str, Union[int, Claim, Evidence]]],\n",
        "    predictions: pd.Series,\n",
        ") -> None:\n",
        "    recall = 0\n",
        "    count = 0\n",
        "\n",
        "    for i, d in enumerate(data):\n",
        "        if d[\"label\"] == \"NOT ENOUGH INFO\":\n",
        "            continue\n",
        "\n",
        "        gt_pages = set([\n",
        "            evidence[2]\n",
        "            for evidence_set in d[\"evidence\"]\n",
        "            for evidence in evidence_set\n",
        "        ])\n",
        "        predicted_pages = predictions.iloc[i]\n",
        "        hits = predicted_pages.intersection(gt_pages)\n",
        "        recall += len(hits) / len(gt_pages)\n",
        "        count += 1\n",
        "\n",
        "    print(f\"Recall: {recall / count}\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The default amount of documents retrieved is at most five documents.  This `num_pred_doc` can be adjusted based on your objective.  Save data in jsonl format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def save_doc(\n",
        "    data: List[Dict[str, Union[int, Claim, Evidence]]],\n",
        "    predictions: pd.Series,\n",
        "    mode: str = \"train\",\n",
        "    num_pred_doc: int = 5,\n",
        ") -> None:\n",
        "    with open(\n",
        "        f\"data/{mode}_doc{num_pred_doc}.jsonl\",\n",
        "        \"w\",\n",
        "        encoding=\"utf8\",\n",
        "    ) as f:\n",
        "        for i, d in enumerate(data):\n",
        "            d[\"predicted_pages\"] = list(predictions.iloc[i])\n",
        "            f.write(json.dumps(d, ensure_ascii=False) + \"\\n\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Main function for document retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ayGI44qkk_wy"
      },
      "outputs": [],
      "source": [
        "def get_pred_pages(series_data: pd.Series) -> Set[Dict[int, str]]:\n",
        "    results = []\n",
        "    tmp_muji = []\n",
        "    # wiki_page: its index showned in claim\n",
        "    mapping = {}\n",
        "    claim = series_data[\"claim\"]\n",
        "    nps = series_data[\"hanlp_results\"]\n",
        "    first_wiki_term = []\n",
        "\n",
        "    for i, np in enumerate(nps):\n",
        "        # Simplified Traditional Chinese Correction\n",
        "        wiki_search_results = [\n",
        "            do_st_corrections(w) for w in wikipedia.search(np)\n",
        "        ]\n",
        "\n",
        "        # Remove the wiki page's description in brackets\n",
        "        wiki_set = [re.sub(r\"\\s\\(\\S+\\)\", \"\", w) for w in wiki_search_results]\n",
        "        wiki_df = pd.DataFrame({\n",
        "            \"wiki_set\": wiki_set,\n",
        "            \"wiki_results\": wiki_search_results\n",
        "        })\n",
        "\n",
        "        # Elements in wiki_set --> index\n",
        "        # Extracting only the first element is one way to avoid extracting\n",
        "        # too many of the similar wiki pages\n",
        "        grouped_df = wiki_df.groupby(\"wiki_set\", sort=False).first()\n",
        "        candidates = grouped_df[\"wiki_results\"].tolist()\n",
        "        # muji refers to wiki_set\n",
        "        muji = grouped_df.index.tolist()\n",
        "\n",
        "        for prefix, term in zip(muji, candidates):\n",
        "            if prefix not in tmp_muji:\n",
        "                matched = False\n",
        "\n",
        "                # Take at least one term from the first noun phrase\n",
        "                if i == 0:\n",
        "                    first_wiki_term.append(term)\n",
        "\n",
        "                # Walrus operator :=\n",
        "                # https://docs.python.org/3/whatsnew/3.8.html#assignment-expressions\n",
        "                # Through these filters, we are trying to figure out if the term\n",
        "                # is within the claim\n",
        "                if (((new_term := term) in claim) or\n",
        "                    ((new_term := term.replace(\"·\", \"\")) in claim) or\n",
        "                    ((new_term := term.split(\" \")[0]) in claim) or\n",
        "                    ((new_term := term.replace(\"-\", \" \")) in claim)):\n",
        "                    matched = True\n",
        "\n",
        "                elif \"·\" in term:\n",
        "                    splitted = term.split(\"·\")\n",
        "                    for split in splitted:\n",
        "                        if (new_term := split) in claim:\n",
        "                            matched = True\n",
        "                            break\n",
        "\n",
        "                if matched:\n",
        "                    # post-processing\n",
        "                    term = term.replace(\" \", \"_\")\n",
        "                    term = term.replace(\"-\", \"\")\n",
        "                    results.append(term)\n",
        "                    mapping[term] = claim.find(new_term)\n",
        "                    tmp_muji.append(new_term)\n",
        "\n",
        "    # 5 is a hyperparameter\n",
        "    if len(results) > 5:\n",
        "        assert -1 not in mapping.values()\n",
        "        results = sorted(mapping, key=mapping.get)[:5]\n",
        "    elif len(results) < 1:\n",
        "        results = first_wiki_term\n",
        "\n",
        "    return set(results)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 1. Get noun phrases from hanlp consituency parsing tree"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Setup [HanLP](https://github.com/hankcs/HanLP) predictor (1 min)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "predictor = (hanlp.pipeline().append(\n",
        "    hanlp.load(\"FINE_ELECTRA_SMALL_ZH\"),\n",
        "    output_key=\"tok\",\n",
        ").append(\n",
        "    hanlp.load(\"CTB9_CON_ELECTRA_SMALL\"),\n",
        "    output_key=\"con\",\n",
        "    input_key=\"tok\",\n",
        "))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We will skip this process which for creating parsing tree when demo on class"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "hanlp_file = f\"data/hanlp_con_results.pkl\"\n",
        "if Path(hanlp_file).exists():\n",
        "    with open(hanlp_file, \"rb\") as f:\n",
        "        hanlp_results = pickle.load(f)\n",
        "else:\n",
        "    hanlp_results = [get_nps_hanlp(predictor, d) for d in TRAIN_DATA]\n",
        "    with open(hanlp_file, \"wb\") as f:\n",
        "        pickle.dump(hanlp_results, f)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Get pages via wiki online api"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "doc_path = f\"data/train_doc5.jsonl\"\n",
        "if Path(doc_path).exists():\n",
        "    with open(doc_path, \"r\", encoding=\"utf8\") as f:\n",
        "        predicted_results = pd.Series([\n",
        "            set(json.loads(line)[\"predicted_pages\"])\n",
        "            for line in f\n",
        "        ])\n",
        "else:\n",
        "    train_df = pd.DataFrame(TRAIN_DATA)\n",
        "    train_df.loc[:, \"hanlp_results\"] = hanlp_results\n",
        "    predicted_results = train_df.parallel_apply(get_pred_pages, axis=1)\n",
        "    save_doc(TRAIN_DATA, predicted_results, mode=\"train\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 2. Calculate our results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "calculate_precision(TRAIN_DATA, predicted_results)\n",
        "calculate_recall(TRAIN_DATA, predicted_results)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 3. Repeat the same process on test set\n",
        "Create parsing tree"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "hanlp_test_file = f\"data/hanlp_con_test_results.pkl\"\n",
        "if Path(hanlp_test_file).exists():\n",
        "    with open(hanlp_test_file, \"rb\") as f:\n",
        "        hanlp_results = pickle.load(f)\n",
        "else:\n",
        "    hanlp_results = [get_nps_hanlp(predictor, d) for d in TEST_DATA]\n",
        "    with open(hanlp_test_file, \"wb\") as f:\n",
        "        pickle.dump(hanlp_results, f)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Get pages via wiki online api"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "test_doc_path = f\"data/test_doc5.jsonl\"\n",
        "if Path(test_doc_path).exists():\n",
        "    with open(test_doc_path, \"r\", encoding=\"utf8\") as f:\n",
        "        test_results = pd.Series(\n",
        "            [set(json.loads(line)[\"predicted_pages\"]) for line in f])\n",
        "else:\n",
        "    test_df = pd.DataFrame(TEST_DATA)\n",
        "    test_df.loc[:, \"hanlp_results\"] = hanlp_results\n",
        "    test_results = test_df.parallel_apply(get_pred_pages, axis=1)\n",
        "    save_doc(TEST_DATA, test_results, mode=\"test\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Ol4zFkSbjgXF"
      },
      "source": [
        "notebook2\n",
        "## PART 2. Sentence retrieval"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import some libs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GlliDsgXjisj"
      },
      "outputs": [],
      "source": [
        "# built-in libs\n",
        "from pathlib import Path\n",
        "from typing import Dict, List, Set, Tuple, Union\n",
        "\n",
        "# third-party libs\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from pandarallel import pandarallel\n",
        "from sklearn.model_selection import train_test_split\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.optim import AdamW\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from transformers import (\n",
        "    AutoModelForSequenceClassification,\n",
        "    AutoTokenizer,\n",
        "    get_scheduler,\n",
        ")\n",
        "\n",
        "from dataset import BERTDataset, Dataset\n",
        "\n",
        "# local libs\n",
        "from utils import (\n",
        "    generate_evidence_to_wiki_pages_mapping,\n",
        "    jsonl_dir_to_df,\n",
        "    load_json,\n",
        "    load_model,\n",
        "    save_checkpoint,\n",
        "    set_lr_scheduler,\n",
        ")\n",
        "\n",
        "pandarallel.initialize(progress_bar=True, verbose=0, nb_workers=10)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Global variable"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J3BBLE3_hlPi"
      },
      "outputs": [],
      "source": [
        "SEED = 42\n",
        "\n",
        "TRAIN_DATA = load_json(\"data/public_train.jsonl\")\n",
        "TEST_DATA = load_json(\"data/public_test.jsonl\")\n",
        "DOC_DATA = load_json(\"data/train_doc5.jsonl\")\n",
        "\n",
        "LABEL2ID: Dict[str, int] = {\n",
        "    \"supports\": 0,\n",
        "    \"refutes\": 1,\n",
        "    \"NOT ENOUGH INFO\": 2,\n",
        "}\n",
        "ID2LABEL: Dict[int, str] = {v: k for k, v in LABEL2ID.items()}\n",
        "\n",
        "_y = [LABEL2ID[data[\"label\"]] for data in TRAIN_DATA]\n",
        "# GT means Ground Truth\n",
        "TRAIN_GT, DEV_GT = train_test_split(\n",
        "    DOC_DATA,\n",
        "    test_size=0.2,\n",
        "    random_state=SEED,\n",
        "    shuffle=True,\n",
        "    stratify=_y,\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Preload wiki database (1 min)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "wiki_pages = jsonl_dir_to_df(\"data/wiki-pages\")\n",
        "mapping = generate_evidence_to_wiki_pages_mapping(wiki_pages)\n",
        "del wiki_pages"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Helper function"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Calculate precision for sentence retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def evidence_macro_precision(\n",
        "    instance: Dict,\n",
        "    top_rows: pd.DataFrame,\n",
        ") -> Tuple[float, float]:\n",
        "    \"\"\"Calculate precision for sentence retrieval\n",
        "    This function is modified from fever-scorer.\n",
        "    https://github.com/sheffieldnlp/fever-scorer/blob/master/src/fever/scorer.py\n",
        "\n",
        "    Args:\n",
        "        instance (dict): a row of the dev set (dev.jsonl) of test set (test.jsonl)\n",
        "        top_rows (pd.DataFrame): our predictions with the top probabilities\n",
        "\n",
        "        IMPORTANT!!!\n",
        "        instance (dict) should have the key of `evidence`.\n",
        "        top_rows (pd.DataFrame) should have a column `predicted_evidence`.\n",
        "\n",
        "    Returns:\n",
        "        Tuple[float, float]:\n",
        "        [1]: relevant and retrieved (numerator of precision)\n",
        "        [2]: retrieved (denominator of precision)\n",
        "    \"\"\"\n",
        "    this_precision = 0.0\n",
        "    this_precision_hits = 0.0\n",
        "\n",
        "    # Return 0, 0 if label is not enough info since not enough info does not\n",
        "    # contain any evidence.\n",
        "    if instance[\"label\"].upper() != \"NOT ENOUGH INFO\":\n",
        "        # e[2] is the page title, e[3] is the sentence index\n",
        "        all_evi = [[e[2], e[3]]\n",
        "                   for eg in instance[\"evidence\"]\n",
        "                   for e in eg\n",
        "                   if e[3] is not None]\n",
        "        claim = instance[\"claim\"]\n",
        "        predicted_evidence = top_rows[top_rows[\"claim\"] ==\n",
        "                                      claim][\"predicted_evidence\"].tolist()\n",
        "\n",
        "        for prediction in predicted_evidence:\n",
        "            if prediction in all_evi:\n",
        "                this_precision += 1.0\n",
        "            this_precision_hits += 1.0\n",
        "\n",
        "        return (this_precision /\n",
        "                this_precision_hits) if this_precision_hits > 0 else 1.0, 1.0\n",
        "\n",
        "    return 0.0, 0.0"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Calculate recall for sentence retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def evidence_macro_recall(\n",
        "    instance: Dict,\n",
        "    top_rows: pd.DataFrame,\n",
        ") -> Tuple[float, float]:\n",
        "    \"\"\"Calculate recall for sentence retrieval\n",
        "    This function is modified from fever-scorer.\n",
        "    https://github.com/sheffieldnlp/fever-scorer/blob/master/src/fever/scorer.py\n",
        "\n",
        "    Args:\n",
        "        instance (dict): a row of the dev set (dev.jsonl) of test set (test.jsonl)\n",
        "        top_rows (pd.DataFrame): our predictions with the top probabilities\n",
        "\n",
        "        IMPORTANT!!!\n",
        "        instance (dict) should have the key of `evidence`.\n",
        "        top_rows (pd.DataFrame) should have a column `predicted_evidence`.\n",
        "\n",
        "    Returns:\n",
        "        Tuple[float, float]:\n",
        "        [1]: relevant and retrieved (numerator of recall)\n",
        "        [2]: relevant (denominator of recall)\n",
        "    \"\"\"\n",
        "    # We only want to score F1/Precision/Recall of recalled evidence for NEI claims\n",
        "    if instance[\"label\"].upper() != \"NOT ENOUGH INFO\":\n",
        "        # If there's no evidence to predict, return 1\n",
        "        if len(instance[\"evidence\"]) == 0 or all(\n",
        "            [len(eg) == 0 for eg in instance]):\n",
        "            return 1.0, 1.0\n",
        "\n",
        "        claim = instance[\"claim\"]\n",
        "\n",
        "        predicted_evidence = top_rows[top_rows[\"claim\"] ==\n",
        "                                      claim][\"predicted_evidence\"].tolist()\n",
        "\n",
        "        for evidence_group in instance[\"evidence\"]:\n",
        "            evidence = [[e[2], e[3]] for e in evidence_group]\n",
        "            if all([item in predicted_evidence for item in evidence]):\n",
        "                # We only want to score complete groups of evidence. Incomplete\n",
        "                # groups are worthless.\n",
        "                return 1.0, 1.0\n",
        "        return 0.0, 1.0\n",
        "    return 0.0, 0.0"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Calculate the scores of sentence retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def evaluate_retrieval(\n",
        "    probs: np.ndarray,\n",
        "    df_evidences: pd.DataFrame,\n",
        "    ground_truths: pd.DataFrame,\n",
        "    top_n: int = 5,\n",
        "    cal_scores: bool = True,\n",
        "    save_name: str = None,\n",
        ") -> Dict[str, float]:\n",
        "    \"\"\"Calculate the scores of sentence retrieval\n",
        "\n",
        "    Args:\n",
        "        probs (np.ndarray): probabilities of the candidate retrieved sentences\n",
        "        df_evidences (pd.DataFrame): the candiate evidence sentences paired with claims\n",
        "        ground_truths (pd.DataFrame): the loaded data of dev.jsonl or test.jsonl\n",
        "        top_n (int, optional): the number of the retrieved sentences. Defaults to 2.\n",
        "\n",
        "    Returns:\n",
        "        Dict[str, float]: F1 score, precision, and recall\n",
        "    \"\"\"\n",
        "    df_evidences[\"prob\"] = probs\n",
        "    top_rows = (\n",
        "        df_evidences.groupby(\"claim\").apply(\n",
        "        lambda x: x.nlargest(top_n, \"prob\"))\n",
        "        .reset_index(drop=True)\n",
        "    )\n",
        "\n",
        "    if cal_scores:\n",
        "        macro_precision = 0\n",
        "        macro_precision_hits = 0\n",
        "        macro_recall = 0\n",
        "        macro_recall_hits = 0\n",
        "\n",
        "        for i, instance in enumerate(ground_truths):\n",
        "            macro_prec = evidence_macro_precision(instance, top_rows)\n",
        "            macro_precision += macro_prec[0]\n",
        "            macro_precision_hits += macro_prec[1]\n",
        "\n",
        "            macro_rec = evidence_macro_recall(instance, top_rows)\n",
        "            macro_recall += macro_rec[0]\n",
        "            macro_recall_hits += macro_rec[1]\n",
        "\n",
        "        pr = (macro_precision /\n",
        "              macro_precision_hits) if macro_precision_hits > 0 else 1.0\n",
        "        rec = (macro_recall /\n",
        "               macro_recall_hits) if macro_recall_hits > 0 else 0.0\n",
        "        f1 = 2.0 * pr * rec / (pr + rec)\n",
        "\n",
        "    if save_name is not None:\n",
        "        # write doc7_sent5 file\n",
        "        with open(f\"data/{save_name}\", \"w\") as f:\n",
        "            for instance in ground_truths:\n",
        "                claim = instance[\"claim\"]\n",
        "                predicted_evidence = top_rows[\n",
        "                    top_rows[\"claim\"] == claim][\"predicted_evidence\"].tolist()\n",
        "                instance[\"predicted_evidence\"] = predicted_evidence\n",
        "                f.write(json.dumps(instance, ensure_ascii=False) + \"\\n\")\n",
        "\n",
        "    if cal_scores:\n",
        "        return {\"F1 score\": f1, \"Precision\": pr, \"Recall\": rec}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Inference script to get probabilites for the candidate evidence sentences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def get_predicted_probs(\n",
        "    model: nn.Module,\n",
        "    dataloader: Dataset,\n",
        "    device: torch.device,\n",
        ") -> np.ndarray:\n",
        "    \"\"\"Inference script to get probabilites for the candidate evidence sentences\n",
        "\n",
        "    Args:\n",
        "        model: the one from HuggingFace Transformers\n",
        "        dataloader: devset or testset in torch dataloader\n",
        "\n",
        "    Returns:\n",
        "        np.ndarray: probabilites of the candidate evidence sentences\n",
        "    \"\"\"\n",
        "    model.eval()\n",
        "    probs = []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        for batch in tqdm(dataloader):\n",
        "            batch = {k: v.to(device) for k, v in batch.items()}\n",
        "            outputs = model(**batch)\n",
        "            logits = outputs.logits\n",
        "            probs.extend(torch.softmax(logits, dim=1)[:, 1].tolist())\n",
        "\n",
        "    return np.array(probs)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "AicupTopkEvidenceBERTDataset class for AICUP dataset with top-k evidence sentences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class SentRetrievalBERTDataset(BERTDataset):\n",
        "    \"\"\"AicupTopkEvidenceBERTDataset class for AICUP dataset with top-k evidence sentences.\"\"\"\n",
        "\n",
        "    def __getitem__(\n",
        "        self,\n",
        "        idx: int,\n",
        "        **kwargs,\n",
        "    ) -> Tuple[Dict[str, torch.Tensor], int]:\n",
        "        item = self.data.iloc[idx]\n",
        "        sentA = item[\"claim\"]\n",
        "        sentB = item[\"text\"]\n",
        "\n",
        "        # claim [SEP] text\n",
        "        concat = self.tokenizer(\n",
        "            sentA,\n",
        "            sentB,\n",
        "            padding=\"max_length\",\n",
        "            max_length=self.max_length,\n",
        "            truncation=True,\n",
        "        )\n",
        "        concat_ten = {k: torch.tensor(v) for k, v in concat.items()}\n",
        "        if \"label\" in item:\n",
        "            concat_ten[\"labels\"] = torch.tensor(item[\"label\"])\n",
        "\n",
        "        return concat_ten"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Main function for sentence retrieval"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gpvXpFwXszfv"
      },
      "outputs": [],
      "source": [
        "def pair_with_wiki_sentences(\n",
        "    mapping: Dict[str, Dict[int, str]],\n",
        "    df: pd.DataFrame,\n",
        "    negative_ratio: float,\n",
        ") -> pd.DataFrame:\n",
        "    \"\"\"Only for creating train sentences.\"\"\"\n",
        "    claims = []\n",
        "    sentences = []\n",
        "    labels = []\n",
        "\n",
        "    # positive\n",
        "    for i in range(len(df)):\n",
        "        if df[\"label\"].iloc[i] == \"NOT ENOUGH INFO\":\n",
        "            continue\n",
        "\n",
        "        claim = df[\"claim\"].iloc[i]\n",
        "        evidence_sets = df[\"evidence\"].iloc[i]\n",
        "        for evidence_set in evidence_sets:\n",
        "            sents = []\n",
        "            for evidence in evidence_set:\n",
        "                # evidence[2] is the page title\n",
        "                page = evidence[2].replace(\" \", \"_\")\n",
        "                # the only page with weird name\n",
        "                if page == \"臺灣海峽危機#第二次臺灣海峽危機（1958）\":\n",
        "                    continue\n",
        "                # evidence[3] is in form of int however, mapping requires str\n",
        "                sent_idx = str(evidence[3])\n",
        "                sents.append(mapping[page][sent_idx])\n",
        "\n",
        "            whole_evidence = \" \".join(sents)\n",
        "\n",
        "            claims.append(claim)\n",
        "            sentences.append(whole_evidence)\n",
        "            labels.append(1)\n",
        "\n",
        "    # negative\n",
        "    for i in range(len(df)):\n",
        "        if df[\"label\"].iloc[i] == \"NOT ENOUGH INFO\":\n",
        "            continue\n",
        "        claim = df[\"claim\"].iloc[i]\n",
        "\n",
        "        evidence_set = set([(evidence[2], evidence[3])\n",
        "                            for evidences in df[\"evidence\"][i]\n",
        "                            for evidence in evidences])\n",
        "        predicted_pages = df[\"predicted_pages\"][i]\n",
        "        for page in predicted_pages:\n",
        "            page = page.replace(\" \", \"_\")\n",
        "            try:\n",
        "                page_sent_id_pairs = [\n",
        "                    (page, sent_idx) for sent_idx in mapping[page].keys()\n",
        "                ]\n",
        "            except KeyError:\n",
        "                # print(f\"{page} is not in our Wiki db.\")\n",
        "                continue\n",
        "\n",
        "            for pair in page_sent_id_pairs:\n",
        "                if pair in evidence_set:\n",
        "                    continue\n",
        "                text = mapping[page][pair[1]]\n",
        "                # `np.random.rand(1) <= 0.05`: Control not to add too many negative samples\n",
        "                if text != \"\" and np.random.rand(1) <= negative_ratio:\n",
        "                    claims.append(claim)\n",
        "                    sentences.append(text)\n",
        "                    labels.append(0)\n",
        "\n",
        "    return pd.DataFrame({\"claim\": claims, \"text\": sentences, \"label\": labels})\n",
        "\n",
        "\n",
        "def pair_with_wiki_sentences_eval(\n",
        "    mapping: Dict[str, Dict[int, str]],\n",
        "    df: pd.DataFrame,\n",
        "    is_testset: bool = False,\n",
        ") -> pd.DataFrame:\n",
        "    \"\"\"Only for creating dev and test sentences.\"\"\"\n",
        "    claims = []\n",
        "    sentences = []\n",
        "    evidence = []\n",
        "    predicted_evidence = []\n",
        "\n",
        "    # negative\n",
        "    for i in range(len(df)):\n",
        "        # if df[\"label\"].iloc[i] == \"NOT ENOUGH INFO\":\n",
        "        #     continue\n",
        "        claim = df[\"claim\"].iloc[i]\n",
        "\n",
        "        predicted_pages = df[\"predicted_pages\"][i]\n",
        "        for page in predicted_pages:\n",
        "            page = page.replace(\" \", \"_\")\n",
        "            try:\n",
        "                page_sent_id_pairs = [(page, k) for k in mapping[page]]\n",
        "            except KeyError:\n",
        "                # print(f\"{page} is not in our Wiki db.\")\n",
        "                continue\n",
        "\n",
        "            for page_name, sentence_id in page_sent_id_pairs:\n",
        "                text = mapping[page][sentence_id]\n",
        "                if text != \"\":\n",
        "                    claims.append(claim)\n",
        "                    sentences.append(text)\n",
        "                    if not is_testset:\n",
        "                        evidence.append(df[\"evidence\"].iloc[i])\n",
        "                    predicted_evidence.append([page_name, int(sentence_id)])\n",
        "\n",
        "    return pd.DataFrame({\n",
        "        \"claim\": claims,\n",
        "        \"text\": sentences,\n",
        "        \"evidence\": evidence if not is_testset else None,\n",
        "        \"predicted_evidence\": predicted_evidence,\n",
        "    })"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 1. Setup training environment"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Hyperparams"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title  { display-mode: \"form\" }\n",
        "\n",
        "MODEL_NAME = \"bert-base-chinese\"  #@param {type:\"string\"}\n",
        "NUM_EPOCHS = 1  #@param {type:\"integer\"}\n",
        "LR = 2e-5  #@param {type:\"number\"}\n",
        "TRAIN_BATCH_SIZE = 64  #@param {type:\"integer\"}\n",
        "TEST_BATCH_SIZE = 256  #@param {type:\"integer\"}\n",
        "NEGATIVE_RATIO = 0.03  #@param {type:\"number\"}\n",
        "VALIDATION_STEP = 50  #@param {type:\"integer\"}\n",
        "TOP_N = 5  #@param {type:\"integer\"}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Experiment Directory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ToLvE9oxIXQo"
      },
      "outputs": [],
      "source": [
        "EXP_DIR = f\"sent_retrieval/e{NUM_EPOCHS}_bs{TRAIN_BATCH_SIZE}_\" + f\"{LR}_neg{NEGATIVE_RATIO}_top{TOP_N}\"\n",
        "LOG_DIR = \"logs/\" + EXP_DIR\n",
        "CKPT_DIR = \"checkpoints/\" + EXP_DIR\n",
        "\n",
        "if not Path(LOG_DIR).exists():\n",
        "    Path(LOG_DIR).mkdir(parents=True)\n",
        "\n",
        "if not Path(CKPT_DIR).exists():\n",
        "    Path(CKPT_DIR).mkdir(parents=True)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 2. Combine claims and evidences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4A5vWEzPiXGF"
      },
      "outputs": [],
      "source": [
        "train_df = pair_with_wiki_sentences(\n",
        "    mapping,\n",
        "    pd.DataFrame(TRAIN_GT),\n",
        "    NEGATIVE_RATIO,\n",
        ")\n",
        "counts = train_df[\"label\"].value_counts()\n",
        "print(\"Now using the following train data with 0 (Negative) and 1 (Positive)\")\n",
        "print(counts)\n",
        "\n",
        "dev_evidences = pair_with_wiki_sentences_eval(mapping, pd.DataFrame(DEV_GT))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 3. Start training"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Dataloader things"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "l48WifjeIGui"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
        "\n",
        "train_dataset = SentRetrievalBERTDataset(train_df, tokenizer=tokenizer)\n",
        "val_dataset = SentRetrievalBERTDataset(dev_evidences, tokenizer=tokenizer)\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    train_dataset,\n",
        "    shuffle=True,\n",
        "    batch_size=TRAIN_BATCH_SIZE,\n",
        ")\n",
        "eval_dataloader = DataLoader(val_dataset, batch_size=TEST_BATCH_SIZE)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Save your memory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "del train_df"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4rl_u0YbeQtY"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\n",
        "    \"cpu\")\n",
        "model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)\n",
        "model.to(device)\n",
        "\n",
        "optimizer = AdamW(model.parameters(), lr=LR)\n",
        "num_training_steps = NUM_EPOCHS * len(train_dataloader)\n",
        "lr_scheduler = set_lr_scheduler(optimizer, num_training_steps)\n",
        "\n",
        "writer = SummaryWriter(LOG_DIR)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Please make sure that you are using gpu when training (5 min)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1AHGaKh1eKmg"
      },
      "outputs": [],
      "source": [
        "progress_bar = tqdm(range(num_training_steps))\n",
        "current_steps = 0\n",
        "\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "    model.train()\n",
        "\n",
        "    for batch in train_dataloader:\n",
        "        batch = {k: v.to(device) for k, v in batch.items()}\n",
        "        outputs = model(**batch)\n",
        "        loss = outputs.loss\n",
        "        loss.backward()\n",
        "\n",
        "        optimizer.step()\n",
        "        lr_scheduler.step()\n",
        "        optimizer.zero_grad()\n",
        "        progress_bar.update(1)\n",
        "        writer.add_scalar(\"training_loss\", loss.item(), current_steps)\n",
        "\n",
        "        y_pred = torch.argmax(outputs.logits, dim=1).tolist()\n",
        "        y_true = batch[\"labels\"].tolist()\n",
        "\n",
        "        current_steps += 1\n",
        "\n",
        "        if current_steps % VALIDATION_STEP == 0 and current_steps > 0:\n",
        "            print(\"Start validation\")\n",
        "            probs = get_predicted_probs(model, eval_dataloader, device)\n",
        "\n",
        "            val_results = evaluate_retrieval(\n",
        "                probs=probs,\n",
        "                df_evidences=dev_evidences,\n",
        "                ground_truths=DEV_GT,\n",
        "                top_n=TOP_N,\n",
        "            )\n",
        "            print(val_results)\n",
        "\n",
        "            # log each metric separately to TensorBoard\n",
        "            for metric_name, metric_value in val_results.items():\n",
        "                writer.add_scalar(\n",
        "                    f\"dev_{metric_name}\",\n",
        "                    metric_value,\n",
        "                    current_steps,\n",
        "                )\n",
        "\n",
        "            save_checkpoint(model, CKPT_DIR, current_steps)\n",
        "\n",
        "print(\"Finished training!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir logs"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Validation part (15 mins)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QS1Ei5DAIO5p"
      },
      "outputs": [],
      "source": [
        "ckpt_name = \"model.50.pt\"  #@param {type:\"string\"}\n",
        "model = load_model(model, ckpt_name, CKPT_DIR)\n",
        "print(\"Start final evaluations and write prediction files.\")\n",
        "\n",
        "train_evidences = pair_with_wiki_sentences_eval(\n",
        "    mapping=mapping,\n",
        "    df=pd.DataFrame(TRAIN_GT),\n",
        ")\n",
        "train_set = SentRetrievalBERTDataset(train_evidences, tokenizer)\n",
        "train_dataloader = DataLoader(train_set, batch_size=TEST_BATCH_SIZE)\n",
        "\n",
        "print(\"Start calculating training scores\")\n",
        "probs = get_predicted_probs(model, train_dataloader, device)\n",
        "train_results = evaluate_retrieval(\n",
        "    probs=probs,\n",
        "    df_evidences=train_evidences,\n",
        "    ground_truths=TRAIN_GT,\n",
        "    top_n=TOP_N,\n",
        "    save_name=f\"train_doc5sent{TOP_N}.jsonl\",\n",
        ")\n",
        "print(f\"Training scores => {train_results}\")\n",
        "\n",
        "print(\"Start validation\")\n",
        "probs = get_predicted_probs(model, eval_dataloader, device)\n",
        "val_results = evaluate_retrieval(\n",
        "    probs=probs,\n",
        "    df_evidences=dev_evidences,\n",
        "    ground_truths=DEV_GT,\n",
        "    top_n=TOP_N,\n",
        "    save_name=f\"dev_doc5sent{TOP_N}.jsonl\",\n",
        ")\n",
        "\n",
        "print(f\"Validation scores => {val_results}\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 4. Check on our test data\n",
        "(5 min)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lVFusJqjmex-"
      },
      "outputs": [],
      "source": [
        "test_data = load_json(\"data/test_doc5.jsonl\")\n",
        "\n",
        "test_evidences = pair_with_wiki_sentences_eval(\n",
        "    mapping,\n",
        "    pd.DataFrame(test_data),\n",
        "    is_testset=True,\n",
        ")\n",
        "test_set = SentRetrievalBERTDataset(test_evidences, tokenizer)\n",
        "test_dataloader = DataLoader(test_set, batch_size=TEST_BATCH_SIZE)\n",
        "\n",
        "print(\"Start predicting the test data\")\n",
        "probs = get_predicted_probs(model, test_dataloader, device)\n",
        "evaluate_retrieval(\n",
        "    probs=probs,\n",
        "    df_evidences=test_evidences,\n",
        "    ground_truths=test_data,\n",
        "    top_n=TOP_N,\n",
        "    cal_scores=False,\n",
        "    save_name=f\"test_doc5sent{TOP_N}.jsonl\",\n",
        ")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "lGzl8a5JteT7"
      },
      "source": [
        "notebook3\n",
        "## PART 3. Claim verification"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "import libs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tgA1vcUyzjlx"
      },
      "outputs": [],
      "source": [
        "import pickle\n",
        "from pathlib import Path\n",
        "from typing import Dict, Tuple\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from pandarallel import pandarallel\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "import torch\n",
        "from sklearn.metrics import accuracy_score\n",
        "from torch.optim import AdamW\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "from transformers import (\n",
        "    AutoModelForSequenceClassification,\n",
        "    AutoTokenizer,\n",
        "    get_scheduler,\n",
        ")\n",
        "\n",
        "from dataset import BERTDataset\n",
        "from utils import (\n",
        "    generate_evidence_to_wiki_pages_mapping,\n",
        "    jsonl_dir_to_df,\n",
        "    load_json,\n",
        "    load_model,\n",
        "    save_checkpoint,\n",
        "    set_lr_scheduler,\n",
        ")\n",
        "\n",
        "pandarallel.initialize(progress_bar=True, verbose=0, nb_workers=4)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Global variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "LABEL2ID: Dict[str, int] = {\n",
        "    \"supports\": 0,\n",
        "    \"refutes\": 1,\n",
        "    \"NOT ENOUGH INFO\": 2,\n",
        "}\n",
        "ID2LABEL: Dict[int, str] = {v: k for k, v in LABEL2ID.items()}\n",
        "\n",
        "TRAIN_DATA = load_json(\"data/train_doc5sent5.jsonl\")\n",
        "DEV_DATA = load_json(\"data/dev_doc5sent5.jsonl\")\n",
        "\n",
        "TRAIN_PKL_FILE = Path(\"data/train_doc5sent5.pkl\")\n",
        "DEV_PKL_FILE = Path(\"data/dev_doc5sent5.pkl\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Preload wiki database (same as part 2.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "wiki_pages = jsonl_dir_to_df(\"data/wiki-pages\")\n",
        "mapping = generate_evidence_to_wiki_pages_mapping(wiki_pages,)\n",
        "del wiki_pages"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Helper function"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "AICUP dataset with top-k evidence sentences."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "class AicupTopkEvidenceBERTDataset(BERTDataset):\n",
        "    \"\"\"AICUP dataset with top-k evidence sentences.\"\"\"\n",
        "\n",
        "    def __getitem__(\n",
        "        self,\n",
        "        idx: int,\n",
        "        **kwargs,\n",
        "    ) -> Tuple[Dict[str, torch.Tensor], int]:\n",
        "        item = self.data.iloc[idx]\n",
        "        claim = item[\"claim\"]\n",
        "        evidence = item[\"evidence_list\"]\n",
        "\n",
        "        # In case there are less than topk evidence sentences\n",
        "        pad = [\"[PAD]\"] * (self.topk - len(evidence))\n",
        "        evidence += pad\n",
        "        concat_claim_evidence = \" [SEP] \".join([*claim, *evidence])\n",
        "\n",
        "        concat = self.tokenizer(\n",
        "            concat_claim_evidence,\n",
        "            padding=\"max_length\",\n",
        "            max_length=self.max_length,\n",
        "            truncation=True,\n",
        "        )\n",
        "        label = LABEL2ID[item[\"label\"]] if \"label\" in item else -1\n",
        "        concat_ten = {k: torch.tensor(v) for k, v in concat.items()}\n",
        "\n",
        "        if \"label\" in item:\n",
        "            concat_ten[\"labels\"] = torch.tensor(label)\n",
        "\n",
        "        return concat_ten"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Evaluation function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def run_evaluation(model: torch.nn.Module, dataloader: DataLoader, device):\n",
        "    model.eval()\n",
        "\n",
        "    loss = 0\n",
        "    y_true = []\n",
        "    y_pred = []\n",
        "    with torch.no_grad():\n",
        "        for batch in tqdm(dataloader):\n",
        "            y_true.extend(batch[\"labels\"].tolist())\n",
        "\n",
        "            batch = {k: v.to(device) for k, v in batch.items()}\n",
        "            outputs = model(**batch)\n",
        "            loss += outputs.loss.item()\n",
        "            logits = outputs.logits\n",
        "            y_pred.extend(torch.argmax(logits, dim=1).tolist())\n",
        "\n",
        "    acc = accuracy_score(y_true, y_pred)\n",
        "\n",
        "    return {\"val_loss\": loss / len(dataloader), \"val_acc\": acc}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def run_predict(model: torch.nn.Module, test_dl: DataLoader, device) -> list:\n",
        "    model.eval()\n",
        "\n",
        "    preds = []\n",
        "    for batch in tqdm(test_dl,\n",
        "                      total=len(test_dl),\n",
        "                      leave=False,\n",
        "                      desc=\"Predicting\"):\n",
        "        batch = {k: v.to(device) for k, v in batch.items()}\n",
        "        pred = model(**batch).logits\n",
        "        pred = torch.argmax(pred, dim=1)\n",
        "        preds.extend(pred.tolist())\n",
        "    return preds"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Main function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def join_with_topk_evidence(\n",
        "    df: pd.DataFrame,\n",
        "    mapping: dict,\n",
        "    mode: str = \"train\",\n",
        "    topk: int = 5,\n",
        ") -> pd.DataFrame:\n",
        "    \"\"\"join_with_topk_evidence join the dataset with topk evidence.\n",
        "\n",
        "    Note:\n",
        "        After extraction, the dataset will be like this:\n",
        "               id     label         claim                           evidence            evidence_list\n",
        "        0    4604  supports       高行健...     [[[3393, 3552, 高行健, 0], [...  [高行健 （ ）江西赣州出...\n",
        "        ..    ...       ...            ...                                ...                     ...\n",
        "        945  2095  supports       美國總...  [[[1879, 2032, 吉米·卡特, 16], [...  [卸任后 ， 卡特積極參與...\n",
        "        停各种战争及人質危機的斡旋工作 ， 反对美国小布什政府攻打伊拉克...\n",
        "\n",
        "        [946 rows x 5 columns]\n",
        "\n",
        "    Args:\n",
        "        df (pd.DataFrame): The dataset with evidence.\n",
        "        wiki_pages (pd.DataFrame): The wiki pages dataframe\n",
        "        topk (int, optional): The topk evidence. Defaults to 5.\n",
        "        cache(Union[Path, str], optional): The cache file path. Defaults to None.\n",
        "            If cache is None, return the result directly.\n",
        "\n",
        "    Returns:\n",
        "        pd.DataFrame: The dataset with topk evidence_list.\n",
        "            The `evidence_list` column will be: List[str]\n",
        "    \"\"\"\n",
        "\n",
        "    # format evidence column to List[List[Tuple[str, str, str, str]]]\n",
        "    if \"evidence\" in df.columns:\n",
        "        df[\"evidence\"] = df[\"evidence\"].parallel_map(\n",
        "            lambda x: [[x]] if not isinstance(x[0], list) else [x]\n",
        "            if not isinstance(x[0][0], list) else x)\n",
        "\n",
        "    print(f\"Extracting evidence_list for the {mode} mode ...\")\n",
        "    if mode == \"eval\":\n",
        "        # extract evidence\n",
        "        df[\"evidence_list\"] = df[\"predicted_evidence\"].parallel_map(lambda x: [\n",
        "            mapping.get(evi_id, {}).get(str(evi_idx), \"\")\n",
        "            for evi_id, evi_idx in x  # for each evidence list\n",
        "        ][:topk] if isinstance(x, list) else [])\n",
        "        print(df[\"evidence_list\"][:5])\n",
        "    else:\n",
        "        # extract evidence\n",
        "        df[\"evidence_list\"] = df[\"evidence\"].parallel_map(lambda x: [\n",
        "            \" \".join([  # join evidence\n",
        "                mapping.get(evi_id, {}).get(str(evi_idx), \"\")\n",
        "                for _, _, evi_id, evi_idx in evi_list\n",
        "            ]) if isinstance(evi_list, list) else \"\"\n",
        "            for evi_list in x  # for each evidence list\n",
        "        ][:1] if isinstance(x, list) else [])\n",
        "\n",
        "    return df"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 1. Setup training environment"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Hyperparams"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "#@title  { display-mode: \"form\" }\n",
        "\n",
        "MODEL_NAME = \"bert-base-chinese\"  #@param {type:\"string\"}\n",
        "TRAIN_BATCH_SIZE = 32  #@param {type:\"integer\"}\n",
        "TEST_BATCH_SIZE = 32  #@param {type:\"integer\"}\n",
        "SEED = 42  #@param {type:\"integer\"}\n",
        "LR = 7e-5  #@param {type:\"number\"}\n",
        "NUM_EPOCHS = 20  #@param {type:\"integer\"}\n",
        "MAX_SEQ_LEN = 256  #@param {type:\"integer\"}\n",
        "EVIDENCE_TOPK = 5  #@param {type:\"integer\"}\n",
        "VALIDATION_STEP = 25  #@param {type:\"integer\"}\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Experiment Directory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "OUTPUT_FILENAME = \"submission.jsonl\"\n",
        "\n",
        "EXP_DIR = f\"claim_verification/e{NUM_EPOCHS}_bs{TRAIN_BATCH_SIZE}_\" + f\"{LR}_top{EVIDENCE_TOPK}\"\n",
        "LOG_DIR = \"logs/\" + EXP_DIR\n",
        "CKPT_DIR = \"checkpoints/\" + EXP_DIR\n",
        "\n",
        "if not Path(LOG_DIR).exists():\n",
        "    Path(LOG_DIR).mkdir(parents=True)\n",
        "\n",
        "if not Path(CKPT_DIR).exists():\n",
        "    Path(CKPT_DIR).mkdir(parents=True)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 2. Concat claim and evidences\n",
        "join topk evidence"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if not TRAIN_PKL_FILE.exists():\n",
        "    train_df = join_with_topk_evidence(\n",
        "        pd.DataFrame(TRAIN_DATA),\n",
        "        mapping,\n",
        "        topk=EVIDENCE_TOPK,\n",
        "    )\n",
        "    train_df.to_pickle(TRAIN_PKL_FILE, protocol=4)\n",
        "else:\n",
        "    with open(TRAIN_PKL_FILE, \"rb\") as f:\n",
        "        train_df = pickle.load(f)\n",
        "\n",
        "if not DEV_PKL_FILE.exists():\n",
        "    dev_df = join_with_topk_evidence(\n",
        "        pd.DataFrame(DEV_DATA),\n",
        "        mapping,\n",
        "        mode=\"eval\",\n",
        "        topk=EVIDENCE_TOPK,\n",
        "    )\n",
        "    dev_df.to_pickle(DEV_PKL_FILE, protocol=4)\n",
        "else:\n",
        "    with open(DEV_PKL_FILE, \"rb\") as f:\n",
        "        dev_df = pickle.load(f)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 3. Training"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Prevent CUDA out of memory"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O0rVk3990DlD"
      },
      "outputs": [],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
        "\n",
        "train_dataset = AicupTopkEvidenceBERTDataset(\n",
        "    train_df,\n",
        "    tokenizer=tokenizer,\n",
        "    max_length=MAX_SEQ_LEN,\n",
        ")\n",
        "val_dataset = AicupTopkEvidenceBERTDataset(\n",
        "    dev_df,\n",
        "    tokenizer=tokenizer,\n",
        "    max_length=MAX_SEQ_LEN,\n",
        ")\n",
        "\n",
        "train_dataloader = DataLoader(\n",
        "    train_dataset,\n",
        "    shuffle=True,\n",
        "    batch_size=TRAIN_BATCH_SIZE,\n",
        ")\n",
        "eval_dataloader = DataLoader(val_dataset, batch_size=TEST_BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CzMgs-Zs3sTN"
      },
      "outputs": [],
      "source": [
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\n",
        "    \"cpu\")\n",
        "model = AutoModelForSequenceClassification.from_pretrained(\n",
        "    MODEL_NAME,\n",
        "    num_labels=len(LABEL2ID),\n",
        ")\n",
        "model.to(device)\n",
        "\n",
        "optimizer = AdamW(model.parameters(), lr=LR)\n",
        "num_training_steps = NUM_EPOCHS * len(train_dataloader)\n",
        "lr_scheduler = set_lr_scheduler(optimizer, num_training_steps)\n",
        "\n",
        "writer = SummaryWriter(LOG_DIR)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Training (30 mins)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_aqMjEek3wmu"
      },
      "outputs": [],
      "source": [
        "progress_bar = tqdm(range(num_training_steps))\n",
        "current_steps = 0\n",
        "\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "    model.train()\n",
        "\n",
        "    for batch in train_dataloader:\n",
        "        batch = {k: v.to(device) for k, v in batch.items()}\n",
        "        outputs = model(**batch)\n",
        "        loss = outputs.loss\n",
        "        loss.backward()\n",
        "\n",
        "        optimizer.step()\n",
        "        lr_scheduler.step()\n",
        "        optimizer.zero_grad()\n",
        "        progress_bar.update(1)\n",
        "        writer.add_scalar(\"training_loss\", loss.item(), current_steps)\n",
        "\n",
        "        y_pred = torch.argmax(outputs.logits, dim=1).tolist()\n",
        "        y_true = batch[\"labels\"].tolist()\n",
        "\n",
        "        current_steps += 1\n",
        "\n",
        "        if current_steps % VALIDATION_STEP == 0 and current_steps > 0:\n",
        "            print(\"Start validation\")\n",
        "            val_results = run_evaluation(model, eval_dataloader, device)\n",
        "\n",
        "            # log each metric separately to TensorBoard\n",
        "            for metric_name, metric_value in val_results.items():\n",
        "                print(f\"{metric_name}: {metric_value}\")\n",
        "                writer.add_scalar(f\"{metric_name}\", metric_value, current_steps)\n",
        "\n",
        "            save_checkpoint(\n",
        "                model,\n",
        "                CKPT_DIR,\n",
        "                current_steps,\n",
        "                mark=f\"val_acc={val_results['val_acc']:.4f}\",\n",
        "            )\n",
        "\n",
        "print(\"Finished training!\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Step 4. Make your submission"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zLkfuoAE49mz"
      },
      "outputs": [],
      "source": [
        "TEST_DATA = load_json(\"data/test_doc5sent5.jsonl\")\n",
        "TEST_PKL_FILE = Path(\"data/test_doc5sent5.pkl\")\n",
        "\n",
        "if not TEST_PKL_FILE.exists():\n",
        "    test_df = join_with_topk_evidence(\n",
        "        pd.DataFrame(TEST_DATA),\n",
        "        mapping,\n",
        "        mode=\"eval\",\n",
        "        topk=EVIDENCE_TOPK,\n",
        "    )\n",
        "    test_df.to_pickle(TEST_PKL_FILE, protocol=4)\n",
        "else:\n",
        "    with open(TEST_PKL_FILE, \"rb\") as f:\n",
        "        test_df = pickle.load(f)\n",
        "\n",
        "test_dataset = AicupTopkEvidenceBERTDataset(\n",
        "    test_df,\n",
        "    tokenizer=tokenizer,\n",
        "    max_length=MAX_SEQ_LEN,\n",
        ")\n",
        "test_dataloader = DataLoader(test_dataset, batch_size=TEST_BATCH_SIZE)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tqIjlht8yCMA"
      },
      "outputs": [],
      "source": [
        "ckpt_name = \"val_acc=0.4259_model.750.pt\"  #@param {type:\"string\"}\n",
        "model = load_model(model, ckpt_name, CKPT_DIR)\n",
        "predicted_label = run_predict(model, test_dataloader, device)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Write files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gl9I3ZWW4pHo"
      },
      "outputs": [],
      "source": [
        "predict_dataset = test_df.copy()\n",
        "predict_dataset[\"predicted_label\"] = list(map(ID2LABEL.get, predicted_label))\n",
        "predict_dataset[[\"id\", \"predicted_label\", \"predicted_evidence\"]].to_json(\n",
        "    OUTPUT_FILENAME,\n",
        "    orient=\"records\",\n",
        "    lines=True,\n",
        "    force_ascii=False,\n",
        ")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "venv",
      "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.5"
    },
    "vscode": {
      "interpreter": {
        "hash": "10286f3c74912972f7d1fdceceee5be5b7c77248e5efe5afcbc6a71f24d230fa"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}