step3_add_evidence.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fec9c5ee-1154-44f1-a22d-00a8eb1a7fa1",
   "metadata": {},
   "source": [
    "# Step 3: Semantic related evidence retrieval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1656c9e5-1056-4f14-92f4-8bb010e31eb5",
   "metadata": {},
   "source": [
    "# Readme\n",
    "*This notebook focusing on utilised an encoder only transformer to find 1 more evidences that are not included in top4 selection from BM25*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b8b61b-a58e-4d2f-950f-11167316a70f",
   "metadata": {},
   "source": [
    "# 1.DataSet Processing\n",
    "(You can add as many code blocks and text blocks as you need. However, YOU SHOULD NOT MODIFY the section title)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a8a776c9-326c-4d91-8bf3-6de04d97d931",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import BatchNormalization, TextVectorization, Embedding, Conv1D, GlobalMaxPooling1D, Dense, Dropout, Lambda, Dot, Reshape, GlobalAveragePooling1D, Flatten\n",
    "from tensorflow.keras import Model, Input, optimizers, layers, metrics\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, Callback\n",
    "from tensorflow.keras import backend as K\n",
    "\n",
    "import spacy\n",
    "import string\n",
    "import time\n",
    "from collections import Counter\n",
    "import string\n",
    "import ast"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e7cddac-13b5-43df-894d-febf59a23548",
   "metadata": {},
   "source": [
    "Load the csv that produced by the step 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6dac1469-a236-48eb-b7ec-39af0f25dd2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_lemmatize = pd.read_csv(\"data_folder/train_k05b085_bm25_top20.csv\")\n",
    "dev_lemmatize = pd.read_csv(\"data_folder/dev_k05b085_bm25_top20.csv\")\n",
    "test_lemmatize = pd.read_csv(\"data_folder/test_k05b085_bm25_top20.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dd2dd60-bafd-4a50-9906-e9cdbabdf7a0",
   "metadata": {},
   "source": [
    "Load the csv that produced by the step 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "65059b92-e58a-40a5-88c0-a91f68cc4c11",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = pd.read_csv(\"data_folder/train.csv\")\n",
    "evidence = pd.read_csv(\"data_folder/evidence.csv\")\n",
    "dev = pd.read_csv(\"data_folder/dev.csv\")\n",
    "test = pd.read_csv(\"data_folder/test.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f654c13c-94f0-491e-8a0d-1e6fe2674d11",
   "metadata": {},
   "source": [
    "Append the top20 evidences retrieved to step 1 dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e4b54d24-c2b4-463e-86cc-01cf16cd0c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "train['prefilter_evidence'] = train_lemmatize['prefilter_evidence']\n",
    "dev['prefilter_evidence'] = dev_lemmatize['prefilter_evidence']\n",
    "test['prefilter_evidence'] = test_lemmatize['prefilter_evidence']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd339c85-ad89-4136-bed6-2c4c2bbf36f8",
   "metadata": {},
   "source": [
    "arrange the evidence list in the structure of list of integer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "669ce2ef-a602-4101-b53a-272acb4d4d21",
   "metadata": {},
   "outputs": [],
   "source": [
    "train['evidences'] = train['evidences'].apply(lambda x: [int(e.split('-')[-1]) for e in eval(x)])\n",
    "dev['evidences'] = dev['evidences'].apply(lambda x: [int(e.split('-')[-1]) for e in eval(x)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "422c5fa1-596c-4605-942d-20ed247d9745",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev['prefilter_evidence'] = dev['prefilter_evidence'].apply(lambda x: ast.literal_eval(x))\n",
    "train['prefilter_evidence'] = train['prefilter_evidence'].apply(lambda x: ast.literal_eval(x))\n",
    "test['prefilter_evidence'] = test['prefilter_evidence'].apply(lambda x: ast.literal_eval(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e81e690b-a0ed-476d-ad0c-5d6dfaf91cf1",
   "metadata": {},
   "source": [
    "# Check peformance of top 4 evidences retrieval before we add 1 more"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "8de3e086-ea1e-4abe-8ad3-c24cdd655e50",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Evidence Recall    = 0.14605048859934852\n",
      "Average Evidence Precision = 0.11197068403908794\n",
      "Average Evidence F-Score   = 0.11874967685228272\n",
      "Average Evidence Recall    = 0.17034632034632033\n",
      "Average Evidence Precision = 0.1266233766233766\n",
      "Average Evidence F-Score   = 0.13530199958771388\n"
     ]
    }
   ],
   "source": [
    "def print_retrival_result(dev):\n",
    "    evidence_recall_scores = []\n",
    "    evidence_precision_scores = []\n",
    "    evidence_fscore_scores = []\n",
    "    \n",
    "    # Iterate over each row in the dev DataFrame to evaluate evidence retrieval\n",
    "    for index, row in dev.iterrows():\n",
    "        true_evidences = set(row['evidences'])    \n",
    "        predicted_evidences = set(row['prefilter_evidence'][:4])\n",
    "        # Initialize counters for correct predictions\n",
    "        evidence_correct = len(true_evidences & predicted_evidences)\n",
    "        \n",
    "        # Calculate recall, precision, and F-score\n",
    "        if len(true_evidences) > 0 and len(predicted_evidences) > 0:\n",
    "            evidence_recall = evidence_correct / len(true_evidences)\n",
    "            evidence_precision = evidence_correct / len(predicted_evidences)\n",
    "            if evidence_recall + evidence_precision > 0:\n",
    "                evidence_fscore = (2 * evidence_precision * evidence_recall) / (evidence_precision + evidence_recall)\n",
    "            else:\n",
    "                evidence_fscore = 0.0\n",
    "        else:\n",
    "            evidence_recall = 0.0\n",
    "            evidence_precision = 0.0\n",
    "            evidence_fscore = 0.0\n",
    "    \n",
    "        # Store the scores\n",
    "        evidence_recall_scores.append(evidence_recall)\n",
    "        evidence_precision_scores.append(evidence_precision)\n",
    "        evidence_fscore_scores.append(evidence_fscore)\n",
    "        \n",
    "    # Calculate mean scores across all instances\n",
    "    mean_recall = np.mean(evidence_recall_scores)\n",
    "    mean_precision = np.mean(evidence_precision_scores)\n",
    "    mean_fscore = np.mean(evidence_fscore_scores)\n",
    "    \n",
    "    # Output the aggregate performance\n",
    "    print(f\"Average Evidence Recall    = {mean_recall}\")\n",
    "    print(f\"Average Evidence Precision = {mean_precision}\")\n",
    "    print(f\"Average Evidence F-Score   = {mean_fscore}\")\n",
    "print_retrival_result(train)\n",
    "print_retrival_result(dev)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb83eff0-3308-4a0d-b393-a39f473449fa",
   "metadata": {},
   "source": [
    "Define top k of evidences you want to keep from BM25 retrieval, that is the search range of the rerank model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb3bb585-627d-41cf-9294-90b71beacf15",
   "metadata": {},
   "source": [
    "For example, the top20 rerank result shown in the report should be change RANK_RANGE to 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "6e84769d-44fe-4493-bd1f-f65b5d7af53e",
   "metadata": {},
   "outputs": [],
   "source": [
    "TOP_K_BM25 = 4\n",
    "RANK_RANGE = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "40c316ef-11a6-469d-8d00-2b0f594da72f",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev['prefilter_evidence'] = dev['prefilter_evidence'].apply(lambda x: x[:RANK_RANGE])\n",
    "train['prefilter_evidence'] = train['prefilter_evidence'].apply(lambda x: x[:RANK_RANGE])\n",
    "test['prefilter_evidence'] = test['prefilter_evidence'].apply(lambda x: x[:RANK_RANGE])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "c7e7a004-7cf0-49c1-baa1-4fbac9bc04cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dev['predict_evidence_bm25'] = dev['prefilter_evidence'].apply(lambda x: x[:TOP_K_BM25])\n",
    "train['predict_evidence_bm25'] = train['prefilter_evidence'].apply(lambda x: x[:TOP_K_BM25])\n",
    "test['predict_evidence_bm25'] = test['prefilter_evidence'].apply(lambda x: x[:TOP_K_BM25])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5cce535-5cf9-4ce0-8ed2-4a7e56edcf4f",
   "metadata": {},
   "source": [
    "Append the true evidences and false evidences to a new column in our search range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "fba104fb-cdee-4f38-a5a6-3ad76970ae3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_evidence(row):\n",
    "    prefilter_list = row['prefilter_evidence']\n",
    "    evidence_list = row['evidences']\n",
    "\n",
    "    # Convert to set for efficient operations\n",
    "    evidence_set = set(evidence_list)\n",
    "\n",
    "    # Create prefilter true evidence list\n",
    "    prefilter_true_evidence = [ev for ev in prefilter_list if ev in evidence_set]\n",
    "\n",
    "    # Update prefilter_evidence by removing the prefilter true evidences from the first k\n",
    "    # and also ensure no element from prefilter true evidence is in the top k of the updated list\n",
    "    updated_prefilter_evidence = [ev for ev in prefilter_list if ev not in prefilter_true_evidence]\n",
    "\n",
    "    return pd.Series([prefilter_true_evidence, updated_prefilter_evidence])\n",
    "train[['prefilter_true_evidence', 'prefilter_false_evidence']] = train.apply(filter_evidence, axis=1)\n",
    "dev[['prefilter_true_evidence', 'prefilter_false_evidence']] = dev.apply(filter_evidence, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8433788-25a3-4ddd-ab43-78ff4424dc5c",
   "metadata": {},
   "source": [
    "Only remove punctuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "8628ca00-68c1-40dd-82d5-831ac0066afb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "def remove_punctuation(input_text):\n",
    "    # Lowercase the input text to standardize it\n",
    "    input_text = input_text.lower()\n",
    "    \n",
    "    # Remove punctuation using a translation table\n",
    "    translator = str.maketrans('', '', string.punctuation)\n",
    "    input_text = input_text.translate(translator)\n",
    "    \n",
    "    # Tokenize the text into words by splitting on whitespace\n",
    "    tokens = input_text.split()\n",
    "    \n",
    "    # Join words back into one string and return\n",
    "    return ' '.join(tokens)\n",
    "\n",
    "# Apply the function to your dataframes\n",
    "train['claim_text'] = train['claim_text'].apply(remove_punctuation)\n",
    "dev['claim_text'] = dev['claim_text'].apply(remove_punctuation)\n",
    "test['claim_text'] = test['claim_text'].apply(remove_punctuation)\n",
    "evidence['evidence_text'] = evidence['evidence_text'].apply(remove_punctuation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35cffa9b-a0c0-49e2-8bdb-b16739bb528c",
   "metadata": {},
   "source": [
    "For each claim, we will concat it with each top k evidences one by one, and we will give a label 0 or 1 to each concatnated text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "41fb7ef3-b6d9-44b1-a56b-7b7af3ebc3d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process each row to create new DataFrame entries\n",
    "def create_new_df(df, train = True):    \n",
    "    new_data = {\n",
    "        'concatenated_text': [],\n",
    "        'evidence_label': [],\n",
    "        'instance_id': [],\n",
    "        'evidence_id': [],\n",
    "        'claim_label': []\n",
    "    }\n",
    "    \n",
    "    for index, row in df.iterrows():\n",
    "        # Process evidences (label 1)\n",
    "        for ev_index in row['evidences'] if train else row['prefilter_true_evidence']:\n",
    "            concatenated_text = \"<CLS> \" + row['claim_text'] + \" <SEP> \" + evidence.iloc[ev_index].evidence_text\n",
    "            new_data['concatenated_text'].append(concatenated_text)\n",
    "            new_data['evidence_label'].append(1)\n",
    "            new_data['instance_id'].append(index)\n",
    "            new_data['evidence_id'].append(ev_index)\n",
    "            new_data['claim_label'].append(row['claim_label'])\n",
    "    \n",
    "        # Process prefilter_evidence (label 0)\n",
    "        for ev_index in row['prefilter_false_evidence']:\n",
    "            concatenated_text = \"<CLS> \" + row['claim_text']  + \" <SEP> \" + evidence.iloc[ev_index].evidence_text\n",
    "            new_data['concatenated_text'].append(concatenated_text)\n",
    "            new_data['evidence_label'].append(0)\n",
    "            new_data['instance_id'].append(index)\n",
    "            new_data['evidence_id'].append(ev_index)\n",
    "            new_data['claim_label'].append('NOT_RELEVANT')\n",
    "    new_df = pd.DataFrame(new_data)\n",
    "    return new_df\n",
    "\n",
    "# Create new DataFrame from processed data\n",
    "train_for_retrival = create_new_df(train)\n",
    "dev_for_retrival = create_new_df(dev, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0de89c54-c868-4862-9484-2c317adfa44a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_test_df(df):    \n",
    "    new_data = {\n",
    "        'concatenated_text': [],\n",
    "        'instance_id': [],\n",
    "        'evidence_id': []\n",
    "    }\n",
    "    \n",
    "    for index, row in df.iterrows():   \n",
    "        evidence_index = row['prefilter_evidence']\n",
    "        for ev_index in evidence_index:\n",
    "            concatenated_text = \"<CLS> \" + row['claim_text']  + \" <SEP> \" + evidence.iloc[ev_index].evidence_text\n",
    "            new_data['concatenated_text'].append(concatenated_text)\n",
    "            new_data['instance_id'].append(index)\n",
    "            new_data['evidence_id'].append(ev_index)\n",
    "    new_df = pd.DataFrame(new_data)\n",
    "    return new_df\n",
    "test_for_retrival = create_test_df(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "74351b00-7f77-44f0-aea6-63d5e4692571",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    11424\n",
       "1     4122\n",
       "Name: evidence_label, dtype: int64"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_for_retrival.evidence_label.value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1502f12-3dcf-4609-9a25-28bb9c1207d9",
   "metadata": {},
   "source": [
    "Imbalance evidence label, since marjority of top evidences are fake evidences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "758f896a-f075-48dd-b488-3a6fe33de97d",
   "metadata": {},
   "source": [
    "Therefore we duplicate the true evidences data for balancing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "4d132a1f-9d4a-4539-b4b8-6320971f5b1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def duplicate_true(df):\n",
    "    filter_label_1 = df[df['evidence_label'] == 1]\n",
    "    \n",
    "    # Duplicate these rows 58 times\n",
    "    duplicated_rows = pd.concat([filter_label_1] * 2, ignore_index=True)\n",
    "    \n",
    "    # Concatenate the duplicated rows back to the original DataFrame\n",
    "    new_df = pd.concat([df, duplicated_rows], ignore_index=True)\n",
    "    \n",
    "    return new_df\n",
    "train_for_retrival_duplicate_true = duplicate_true(train_for_retrival)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "4d71bba3-af3a-40ba-b22a-735b9e31ce43",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start processing!\n",
      "Vocabulary size on all claim text: 19141\n",
      "Processed and encoded data.\n"
     ]
    }
   ],
   "source": [
    "# Define file paths\n",
    "sequence_length = 96\n",
    "print(\"Start processing!\")\n",
    "# Concatenate priority texts\n",
    "priority_texts = pd.concat([train_for_retrival['concatenated_text'], dev_for_retrival['concatenated_text'], test_for_retrival['concatenated_text']])\n",
    "\n",
    "# Create the TextVectorization layer\n",
    "vectorize_layer = tf.keras.layers.TextVectorization(\n",
    "    output_mode=\"int\",\n",
    "    output_sequence_length=sequence_length\n",
    ")\n",
    "\n",
    "# Adapt the vectorization layer on priority texts first\n",
    "vectorize_layer.adapt(priority_texts)\n",
    "print(\"Vocabulary size on all claim text:\", len(vectorize_layer.get_vocabulary()))\n",
    "max_features = len(vectorize_layer.get_vocabulary())\n",
    "# Encode texts\n",
    "train_encoded = vectorize_layer(train_for_retrival_duplicate_true['concatenated_text'].to_numpy())\n",
    "dev_encoded = vectorize_layer(dev_for_retrival['concatenated_text'].to_numpy())\n",
    "print(\"Processed and encoded data.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "66feedc8-d97b-4edd-af39-4dab2b608114",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_dataset(encoded_data, df):\n",
    "    features = tf.convert_to_tensor(encoded_data)\n",
    "    evidence_labels = tf.convert_to_tensor(df['evidence_label'].values, dtype=tf.float32)\n",
    "    \n",
    "    # Combine features and labels into a dataset\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((features, {\"evidence_label\": evidence_labels}))\n",
    "    \n",
    "    # Batch the dataset\n",
    "    dataset = dataset.batch(32).shuffle(buffer_size=len(df))\n",
    "    return dataset\n",
    "train_dataset = prepare_dataset(train_encoded, train_for_retrival_duplicate_true)\n",
    "dev_dataset = prepare_dataset(dev_encoded, dev_for_retrival)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a06875d7-1bd7-4913-9d1f-60118874ed08",
   "metadata": {},
   "source": [
    "# 2. Model Implementation\n",
    "(You can add as many code blocks and text blocks as you need. However, YOU SHOULD NOT MODIFY the section title)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adc774c8-3bd9-4ac1-8e6a-e2eee5958847",
   "metadata": {},
   "source": [
    "Self implemented metrics for tracking the TP and TN rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "51fe90ae-6241-4c83-b12f-194152c466b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Label1Accuracy(tf.keras.metrics.Metric):\n",
    "    def __init__(self, name='label_1_acc', **kwargs):\n",
    "        super(Label1Accuracy, self).__init__(name=name, **kwargs)\n",
    "        self.true_positives = self.add_weight(name='tp', initializer='zeros')\n",
    "        self.total_label_1 = self.add_weight(name='total', initializer='zeros')\n",
    "\n",
    "    def update_state(self, y_true, y_pred, sample_weight=None):\n",
    "        y_true = tf.cast(y_true, tf.bool)\n",
    "        y_pred = tf.cast(tf.round(y_pred), tf.bool)\n",
    "        values = tf.logical_and(tf.equal(y_true, True), tf.equal(y_pred, True))\n",
    "        self.true_positives.assign_add(tf.reduce_sum(tf.cast(values, self.dtype)))\n",
    "        self.total_label_1.assign_add(tf.reduce_sum(tf.cast(y_true, self.dtype)))\n",
    "\n",
    "    def result(self):\n",
    "        return self.true_positives / (self.total_label_1 + tf.keras.backend.epsilon())\n",
    "\n",
    "    def reset_states(self):\n",
    "        self.true_positives.assign(0)\n",
    "        self.total_label_1.assign(0)\n",
    "\n",
    "class Label0Accuracy(tf.keras.metrics.Metric):\n",
    "    def __init__(self, name='label_0_acc', **kwargs):\n",
    "        super(Label0Accuracy, self).__init__(name=name, **kwargs)\n",
    "        self.true_negatives = self.add_weight(name='tn', initializer='zeros')\n",
    "        self.total_label_0 = self.add_weight(name='total', initializer='zeros')\n",
    "\n",
    "    def update_state(self, y_true, y_pred, sample_weight=None):\n",
    "        y_true = tf.cast(y_true, tf.bool)\n",
    "        y_pred = tf.cast(tf.round(y_pred), tf.bool)\n",
    "        values = tf.logical_and(tf.equal(y_true, False), tf.equal(y_pred, False))\n",
    "        self.true_negatives.assign_add(tf.reduce_sum(tf.cast(values, self.dtype)))\n",
    "        self.total_label_0.assign_add(tf.reduce_sum(tf.cast(tf.logical_not(y_true), self.dtype)))\n",
    "\n",
    "    def result(self):\n",
    "        return self.true_negatives / (self.total_label_0 + tf.keras.backend.epsilon())\n",
    "\n",
    "    def reset_states(self):\n",
    "        self.true_negatives.assign(0)\n",
    "        self.total_label_0.assign(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "7b7c9648-ccff-4774-a738-9f97d479e905",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_8\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_9 (InputLayer)         [(None, 96)]              0         \n",
      "_________________________________________________________________\n",
      "embedding_16 (Embedding)     (None, 96, 1024)          19600384  \n",
      "_________________________________________________________________\n",
      "tf.__operators__.add_8 (TFOp (None, 96, 1024)          0         \n",
      "_________________________________________________________________\n",
      "transformer_block_16 (Transf (None, 96, 1024)          17057920  \n",
      "_________________________________________________________________\n",
      "transformer_block_17 (Transf (None, 96, 1024)          17057920  \n",
      "_________________________________________________________________\n",
      "global_max_pooling1d_8 (Glob (None, 1024)              0         \n",
      "_________________________________________________________________\n",
      "dropout_44 (Dropout)         (None, 1024)              0         \n",
      "_________________________________________________________________\n",
      "evidence_label (Dense)       (None, 1)                 1025      \n",
      "=================================================================\n",
      "Total params: 53,717,249\n",
      "Trainable params: 53,717,249\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "class TransformerBlock(tf.keras.layers.Layer):\n",
    "    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1, **kwargs):\n",
    "        super(TransformerBlock, self).__init__(**kwargs)  # Pass any extra arguments to the superclass\n",
    "        self.embed_dim = embed_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.ff_dim = ff_dim\n",
    "        self.rate = rate\n",
    "        self.att = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)\n",
    "        self.ffn = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(ff_dim, activation=\"relu\"),\n",
    "            tf.keras.layers.Dense(embed_dim),\n",
    "        ])\n",
    "        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)\n",
    "        self.dropout1 = tf.keras.layers.Dropout(rate)\n",
    "        self.dropout2 = tf.keras.layers.Dropout(rate)\n",
    "\n",
    "    def call(self, inputs, training=False):\n",
    "        attn_output = self.att(inputs, inputs)\n",
    "        attn_output = self.dropout1(attn_output, training=training)\n",
    "        out1 = self.layernorm1(inputs + attn_output)\n",
    "        ffn_output = self.ffn(out1)\n",
    "        ffn_output = self.dropout2(ffn_output, training=training)\n",
    "        return self.layernorm2(out1 + ffn_output)\n",
    "\n",
    "    def get_config(self):\n",
    "        config = super(TransformerBlock, self).get_config()\n",
    "        config.update({\n",
    "            \"embed_dim\": self.embed_dim,\n",
    "            \"num_heads\": self.num_heads,\n",
    "            \"ff_dim\": self.ff_dim,\n",
    "            \"rate\": self.rate\n",
    "        })\n",
    "        return config\n",
    "\n",
    "def create_embedding_model(sequence_length, max_features, num_transformer_blocks, embedding_dim, num_heads, ff_dim, rate):\n",
    "    input_layer = Input(shape=(sequence_length,), dtype=\"int64\")\n",
    "    \n",
    "    # Embedding layer with positional encoding\n",
    "    embedding_layer = Embedding(max_features, embedding_dim)\n",
    "    x = embedding_layer(input_layer)\n",
    "    # Adding positional encoding\n",
    "    \n",
    "    position_embedding = Embedding(input_dim=sequence_length, output_dim=embedding_dim)\n",
    "    positions = tf.range(start=0, limit=sequence_length, delta=1)\n",
    "    x += position_embedding(positions)\n",
    "    \n",
    "    # Transformer blocks\n",
    "    for _ in range(num_transformer_blocks):\n",
    "        x = TransformerBlock(embedding_dim, num_heads, ff_dim, rate)(x)\n",
    "    x = GlobalMaxPooling1D()(x)\n",
    "    x = Dropout(0.1)(x)\n",
    "    # Dense layers and output\n",
    "    evidence_output = Dense(1, activation='sigmoid', name=\"evidence_label\")(x)\n",
    "    \n",
    "    # Compile model\n",
    "    model = Model(inputs=input_layer, outputs=evidence_output)\n",
    "    return model\n",
    "\n",
    "# Initialize and compile the model as before\n",
    "embedding_model = create_embedding_model(sequence_length=sequence_length, max_features=max_features, num_transformer_blocks = 2,\n",
    "                                         embedding_dim = 1024, num_heads=4, ff_dim=128, rate=0.1)\n",
    "embedding_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "787ea040-9c7a-4726-b6dd-723b6c02a545",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_model.compile(\n",
    "    optimizer=optimizers.Adam(1e-5),\n",
    "    loss={\n",
    "        'evidence_label': 'binary_crossentropy'\n",
    "    },\n",
    "    metrics={\n",
    "        'evidence_label': ['accuracy', Label1Accuracy(), Label0Accuracy()]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "369db270-1fc3-4091-b2c3-2776f95a5785",
   "metadata": {},
   "source": [
    "Typically 10-20 epochs could get the following result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "74d1bf41-9ee0-4076-9a0b-3b0fc760c278",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "744/744 [==============================] - 29s 38ms/step - loss: 0.8569 - accuracy: 0.5082 - label_1_acc: 0.5273 - label_0_acc: 0.4876 - val_loss: 0.8926 - val_accuracy: 0.2519 - val_label_1_acc: 0.7542 - val_label_0_acc: 0.2103\n",
      "Epoch 2/10\n",
      "744/744 [==============================] - 28s 37ms/step - loss: 0.7674 - accuracy: 0.5439 - label_1_acc: 0.5716 - label_0_acc: 0.5139 - val_loss: 0.9755 - val_accuracy: 0.1429 - val_label_1_acc: 0.8983 - val_label_0_acc: 0.0802\n",
      "Epoch 3/10\n",
      "744/744 [==============================] - 28s 38ms/step - loss: 0.7179 - accuracy: 0.5804 - label_1_acc: 0.6040 - label_0_acc: 0.5549 - val_loss: 0.7057 - val_accuracy: 0.5227 - val_label_1_acc: 0.5085 - val_label_0_acc: 0.5239\n",
      "Epoch 4/10\n",
      "744/744 [==============================] - 28s 37ms/step - loss: 0.6797 - accuracy: 0.6053 - label_1_acc: 0.6301 - label_0_acc: 0.5784 - val_loss: 0.7572 - val_accuracy: 0.4695 - val_label_1_acc: 0.5763 - val_label_0_acc: 0.4606\n",
      "Epoch 5/10\n",
      "744/744 [==============================] - 28s 37ms/step - loss: 0.6274 - accuracy: 0.6544 - label_1_acc: 0.6809 - label_0_acc: 0.6258 - val_loss: 0.7785 - val_accuracy: 0.4610 - val_label_1_acc: 0.6186 - val_label_0_acc: 0.4480\n",
      "Epoch 6/10\n",
      "744/744 [==============================] - 28s 37ms/step - loss: 0.5869 - accuracy: 0.6844 - label_1_acc: 0.7014 - label_0_acc: 0.6659 - val_loss: 1.0375 - val_accuracy: 0.2221 - val_label_1_acc: 0.9153 - val_label_0_acc: 0.1646\n",
      "Epoch 7/10\n",
      "744/744 [==============================] - 28s 38ms/step - loss: 0.5499 - accuracy: 0.7179 - label_1_acc: 0.7378 - label_0_acc: 0.6963 - val_loss: 0.8137 - val_accuracy: 0.4500 - val_label_1_acc: 0.7627 - val_label_0_acc: 0.4241\n",
      "Epoch 8/10\n",
      "744/744 [==============================] - 28s 38ms/step - loss: 0.4941 - accuracy: 0.7606 - label_1_acc: 0.7822 - label_0_acc: 0.7372 - val_loss: 0.4842 - val_accuracy: 0.7883 - val_label_1_acc: 0.2373 - val_label_0_acc: 0.8340\n",
      "Epoch 9/10\n",
      "744/744 [==============================] - 28s 37ms/step - loss: 0.4431 - accuracy: 0.7955 - label_1_acc: 0.8134 - label_0_acc: 0.7762 - val_loss: 0.7803 - val_accuracy: 0.5214 - val_label_1_acc: 0.6780 - val_label_0_acc: 0.5084\n",
      "Epoch 10/10\n",
      "744/744 [==============================] - 28s 38ms/step - loss: 0.3869 - accuracy: 0.8313 - label_1_acc: 0.8557 - label_0_acc: 0.8048 - val_loss: 0.5420 - val_accuracy: 0.7201 - val_label_1_acc: 0.4068 - val_label_0_acc: 0.7461\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x168459fc160>"
      ]
     },
     "execution_count": 94,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_model.fit(\n",
    "    train_dataset,\n",
    "    validation_data=dev_dataset,\n",
    "    epochs=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42e66843-6d1a-46ff-962a-a8cda6e694ed",
   "metadata": {},
   "source": [
    "# 3.Testing and Evaluation\n",
    "(You can add as many code blocks and text blocks as you need. However, YOU SHOULD NOT MODIFY the section title)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e2b54ce-13b5-4e89-894f-2d75d9440ed3",
   "metadata": {},
   "source": [
    "# Valid actual performence of dev data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "8bb08a15-f2e8-4688-b665-37a8451ac4fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def attached_transformer_prediction(df, df_retrival, embedding_model, text_encoded, top_add_k=5, top_remove_k = 5):\n",
    "    # Predict the probabilities for each text\n",
    "    prob_predictions = embedding_model.predict(text_encoded)\n",
    "    df_retrival['predicted_label'] = prob_predictions.flatten()\n",
    "    \n",
    "    # Get the top k evidence by probability for each instance\n",
    "    top_true_evidences = df_retrival.sort_values(by=['instance_id', 'predicted_label'], ascending=[True, False])\n",
    "    top_true_evidences = top_true_evidences.groupby('instance_id').head(top_add_k)\n",
    "\n",
    "    top_false_evidences = df_retrival.sort_values(by=['instance_id', 'predicted_label'], ascending=[True, True])\n",
    "    top_false_evidences = top_false_evidences.groupby('instance_id').head(top_remove_k)\n",
    "    \n",
    "    # Aggregate the top k evidences and their probabilities into lists\n",
    "    aggregated_add_evidences = top_true_evidences.groupby('instance_id')['evidence_id'].apply(list).reset_index(name='predicted_add_evidence_transformer')\n",
    "    aggregated_remove_evidences = top_false_evidences.groupby('instance_id')['evidence_id'].apply(list).reset_index(name='predicted_rm_evidence_transformer')\n",
    "    \n",
    "    # Merge the aggregated data back into the original DataFrame using the index (instance_id)\n",
    "    df = df.merge(aggregated_add_evidences, how='left', left_index=True, right_on='instance_id')\n",
    "    df = df.merge(aggregated_remove_evidences, how='left', left_index=True, right_on='instance_id')\n",
    "\n",
    "    # Cleanup the DataFrame to remove any temporary columns or duplicate index columns\n",
    "    df.drop(columns=['instance_id_x', 'instance_id_y'], errors='ignore', inplace=True)\n",
    "    \n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "0def5f63-7cd9-4e88-8e92-162f8ca7f216",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_with_result = attached_transformer_prediction(test, test_for_retrival , embedding_model, \n",
    "                                                   vectorize_layer(test_for_retrival['concatenated_text'].to_numpy()))\n",
    "dev_with_result = attached_transformer_prediction(dev, dev_for_retrival , embedding_model, \n",
    "                                                   dev_encoded)\n",
    "train_with_result = attached_transformer_prediction(train, train_for_retrival , embedding_model, \n",
    "                                                   vectorize_layer(train_for_retrival['concatenated_text'].to_numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e5b090b-e8f0-47ec-a04f-83813cddf553",
   "metadata": {},
   "source": [
    "top 4 fully rerank by transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "319bdcd6-f742-41d8-99e6-dd0876789948",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Evidence Recall    = 0.14653679653679655\n",
      "Average Evidence Precision = 0.1038961038961039\n",
      "Average Evidence F-Score   = 0.11231704803133374\n"
     ]
    }
   ],
   "source": [
    "evidence_recall_scores = []\n",
    "evidence_precision_scores = []\n",
    "evidence_fscore_scores = []\n",
    "\n",
    "# Iterate over each row in the dev DataFrame to evaluate evidence retrieval\n",
    "for index, row in dev_with_result.iterrows():\n",
    "    true_evidences = set(row['evidences'])\n",
    "    # Create copies of the lists to avoid modifying the original DataFrame\n",
    "    # Combine the filtered transformer evidences with BM25 predictions\n",
    "    predicted_evidences = []\n",
    "    for ev in row['predicted_add_evidence_transformer'][:4]:\n",
    "        if ev not in predicted_evidences:\n",
    "            predicted_evidences.append(ev)\n",
    "    # Initialize counters for correct predictions\n",
    "    evidence_correct = len(true_evidences & set(predicted_evidences))\n",
    "    \n",
    "    # Calculate recall, precision, and F-score\n",
    "    if len(true_evidences) > 0 and len(predicted_evidences) > 0:\n",
    "        evidence_recall = evidence_correct / len(true_evidences)\n",
    "        evidence_precision = evidence_correct / len(predicted_evidences)\n",
    "        if evidence_recall + evidence_precision > 0:\n",
    "            evidence_fscore = (2 * evidence_precision * evidence_recall) / (evidence_precision + evidence_recall)\n",
    "        else:\n",
    "            evidence_fscore = 0.0\n",
    "    else:\n",
    "        evidence_recall = 0.0\n",
    "        evidence_precision = 0.0\n",
    "        evidence_fscore = 0.0\n",
    "\n",
    "    # Store the scores\n",
    "    evidence_recall_scores.append(evidence_recall)\n",
    "    evidence_precision_scores.append(evidence_precision)\n",
    "    evidence_fscore_scores.append(evidence_fscore)\n",
    "    \n",
    "# Calculate mean scores across all instances\n",
    "mean_recall = np.mean(evidence_recall_scores)\n",
    "mean_precision = np.mean(evidence_precision_scores)\n",
    "mean_fscore = np.mean(evidence_fscore_scores)\n",
    "\n",
    "# Output the aggregate performance\n",
    "print(f\"Average Evidence Recall    = {mean_recall}\")\n",
    "print(f\"Average Evidence Precision = {mean_precision}\")\n",
    "print(f\"Average Evidence F-Score   = {mean_fscore}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cde4fba-6a38-4237-9604-2e02c134e27e",
   "metadata": {},
   "source": [
    "top 1 rerank model recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "707f6991-e42f-4710-a7d1-73687f3db37c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Evidence Recall    = 0.039718614718614716\n",
      "Average Evidence Precision = 0.12337662337662338\n",
      "Average Evidence F-Score   = 0.05670995670995672\n"
     ]
    }
   ],
   "source": [
    "evidence_recall_scores = []\n",
    "evidence_precision_scores = []\n",
    "evidence_fscore_scores = []\n",
    "\n",
    "# Iterate over each row in the dev DataFrame to evaluate evidence retrieval\n",
    "for index, row in dev_with_result.iterrows():\n",
    "    true_evidences = set(row['evidences'])\n",
    "    # Create copies of the lists to avoid modifying the original DataFrame\n",
    "    # Combine the filtered transformer evidences with BM25 predictions\n",
    "    predicted_evidences = []\n",
    "    for ev in row['predicted_add_evidence_transformer'][:1]:\n",
    "        if ev not in predicted_evidences:\n",
    "            predicted_evidences.append(ev)\n",
    "    # Initialize counters for correct predictions\n",
    "    evidence_correct = len(true_evidences & set(predicted_evidences))\n",
    "    \n",
    "    # Calculate recall, precision, and F-score\n",
    "    if len(true_evidences) > 0 and len(predicted_evidences) > 0:\n",
    "        evidence_recall = evidence_correct / len(true_evidences)\n",
    "        evidence_precision = evidence_correct / len(predicted_evidences)\n",
    "        if evidence_recall + evidence_precision > 0:\n",
    "            evidence_fscore = (2 * evidence_precision * evidence_recall) / (evidence_precision + evidence_recall)\n",
    "        else:\n",
    "            evidence_fscore = 0.0\n",
    "    else:\n",
    "        evidence_recall = 0.0\n",
    "        evidence_precision = 0.0\n",
    "        evidence_fscore = 0.0\n",
    "\n",
    "    # Store the scores\n",
    "    evidence_recall_scores.append(evidence_recall)\n",
    "    evidence_precision_scores.append(evidence_precision)\n",
    "    evidence_fscore_scores.append(evidence_fscore)\n",
    "    \n",
    "# Calculate mean scores across all instances\n",
    "mean_recall = np.mean(evidence_recall_scores)\n",
    "mean_precision = np.mean(evidence_precision_scores)\n",
    "mean_fscore = np.mean(evidence_fscore_scores)\n",
    "\n",
    "# Output the aggregate performance\n",
    "print(f\"Average Evidence Recall    = {mean_recall}\")\n",
    "print(f\"Average Evidence Precision = {mean_precision}\")\n",
    "print(f\"Average Evidence F-Score   = {mean_fscore}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bf9f674-84e4-4361-b91c-e44bd49c0e29",
   "metadata": {},
   "source": [
    "# Final ensemble result of top 4 BM25 and top 1 rerank model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 462,
   "id": "1e0eeb51-f697-412e-9735-c82931a5fd0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average Evidence Recall    = 0.20086580086580083\n",
      "Average Evidence Precision = 0.13084415584415585\n",
      "Average Evidence F-Score   = 0.1473768295196867\n"
     ]
    }
   ],
   "source": [
    "evidence_recall_scores = []\n",
    "evidence_precision_scores = []\n",
    "evidence_fscore_scores = []\n",
    "\n",
    "# Iterate over each row in the dev DataFrame to evaluate evidence retrieval\n",
    "for index, row in dev_with_result.iterrows():\n",
    "    true_evidences = set(row['evidences'])\n",
    "    # Create copies of the lists to avoid modifying the original DataFrame\n",
    "    # Combine the filtered transformer evidences with BM25 predictions\n",
    "    predicted_evidences = []\n",
    "    for ev in row['predict_evidence_bm25']:\n",
    "        #if ev not in row['predicted_rm_evidence_transformer'][:1]:\n",
    "        predicted_evidences.append(ev)\n",
    "    for ev in row['predicted_add_evidence_transformer'][:1]:\n",
    "        if ev not in predicted_evidences:\n",
    "            predicted_evidences.append(ev)\n",
    "    # Initialize counters for correct predictions\n",
    "    evidence_correct = len(true_evidences & set(predicted_evidences))\n",
    "    \n",
    "    # Calculate recall, precision, and F-score\n",
    "    if len(true_evidences) > 0 and len(predicted_evidences) > 0:\n",
    "        evidence_recall = evidence_correct / len(true_evidences)\n",
    "        evidence_precision = evidence_correct / len(predicted_evidences)\n",
    "        if evidence_recall + evidence_precision > 0:\n",
    "            evidence_fscore = (2 * evidence_precision * evidence_recall) / (evidence_precision + evidence_recall)\n",
    "        else:\n",
    "            evidence_fscore = 0.0\n",
    "    else:\n",
    "        evidence_recall = 0.0\n",
    "        evidence_precision = 0.0\n",
    "        evidence_fscore = 0.0\n",
    "\n",
    "    # Store the scores\n",
    "    evidence_recall_scores.append(evidence_recall)\n",
    "    evidence_precision_scores.append(evidence_precision)\n",
    "    evidence_fscore_scores.append(evidence_fscore)\n",
    "    \n",
    "# Calculate mean scores across all instances\n",
    "mean_recall = np.mean(evidence_recall_scores)\n",
    "mean_precision = np.mean(evidence_precision_scores)\n",
    "mean_fscore = np.mean(evidence_fscore_scores)\n",
    "\n",
    "# Output the aggregate performance\n",
    "print(f\"Average Evidence Recall    = {mean_recall}\")\n",
    "print(f\"Average Evidence Precision = {mean_precision}\")\n",
    "print(f\"Average Evidence F-Score   = {mean_fscore}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa0b1db4-caec-4f55-9771-db5af02c61b1",
   "metadata": {},
   "source": [
    "The F-score imporve from 0.135 to 0.147"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5056adf-5a4d-4e9e-af96-69f869da2bbe",
   "metadata": {},
   "source": [
    "# The step 4 notebook will use the data that contains 1 more evidences added by transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 457,
   "id": "d6e8ffdd-1791-41a4-b97e-137f7e8ee6c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_with_result.to_csv(\"train_retrival_result_with_transformer.csv\")\n",
    "dev_with_result.to_csv(\"dev_retrival_result_with_transformer.csv\")\n",
    "test_with_result.to_csv(\"test_retrival_result_with_transformer.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae4f475-4e47-4f48-9a78-4a4097b9d70e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create test file"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}