2_MU-Kindai_SBERT-JSNLI-base.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6120c9f4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.949996Z",
     "start_time": "2023-10-07T08:50:44.497041Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:46.918794Z",
     "iopub.status.busy": "2023-10-08T08:43:46.917790Z",
     "iopub.status.idle": "2023-10-08T08:43:48.302919Z",
     "shell.execute_reply": "2023-10-08T08:43:48.302919Z"
    },
    "papermill": {
     "duration": 1.391007,
     "end_time": "2023-10-08T08:43:48.304922",
     "exception": false,
     "start_time": "2023-10-08T08:43:46.913915",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cebf55d5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.965421Z",
     "start_time": "2023-10-07T08:50:45.951998Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:48.313434Z",
     "iopub.status.busy": "2023-10-08T08:43:48.312428Z",
     "iopub.status.idle": "2023-10-08T08:43:48.318536Z",
     "shell.execute_reply": "2023-10-08T08:43:48.318536Z"
    },
    "papermill": {
     "duration": 0.012577,
     "end_time": "2023-10-08T08:43:48.320497",
     "exception": false,
     "start_time": "2023-10-08T08:43:48.307920",
     "status": "completed"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "jsts_url = \"https://raw.githubusercontent.com/yahoojapan/JGLUE/main/datasets/jsts-v1.1/valid-v1.1.json\"\n",
    "jsick_url = \"https://github.com/verypluming/JSICK/raw/main/jsick/test.tsv\"\n",
    "miracle_n_hard_negs = 300\n",
    "miracle_n_recall = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "25b12dde",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:48.328036Z",
     "iopub.status.busy": "2023-10-08T08:43:48.328036Z",
     "iopub.status.idle": "2023-10-08T08:43:48.334175Z",
     "shell.execute_reply": "2023-10-08T08:43:48.334068Z"
    },
    "papermill": {
     "duration": 0.012056,
     "end_time": "2023-10-08T08:43:48.336089",
     "exception": false,
     "start_time": "2023-10-08T08:43:48.324033",
     "status": "completed"
    },
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "model_id = \"MU-Kindai/SBERT-JSNLI-base\"\n",
    "query_prefix = \"\"\n",
    "passage_prefix = \"\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05495700",
   "metadata": {
    "papermill": {
     "duration": 0.002527,
     "end_time": "2023-10-08T08:43:48.341974",
     "exception": false,
     "start_time": "2023-10-08T08:43:48.339447",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9676eb79",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:53.766453Z",
     "start_time": "2023-10-07T08:50:51.631864Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:48.349996Z",
     "iopub.status.busy": "2023-10-08T08:43:48.349996Z",
     "iopub.status.idle": "2023-10-08T08:43:50.333582Z",
     "shell.execute_reply": "2023-10-08T08:43:50.332582Z"
    },
    "papermill": {
     "duration": 1.990109,
     "end_time": "2023-10-08T08:43:50.335090",
     "exception": false,
     "start_time": "2023-10-08T08:43:48.344981",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer(model_id)\n",
    "model.max_seq_length = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "874638a3",
   "metadata": {
    "papermill": {
     "duration": 0.003508,
     "end_time": "2023-10-08T08:43:50.343604",
     "exception": false,
     "start_time": "2023-10-08T08:43:50.340096",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77838b50",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.215161Z",
     "start_time": "2023-10-07T08:50:53.767452Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:50.355117Z",
     "iopub.status.busy": "2023-10-08T08:43:50.354119Z",
     "iopub.status.idle": "2023-10-08T08:43:51.052130Z",
     "shell.execute_reply": "2023-10-08T08:43:51.051627Z"
    },
    "papermill": {
     "duration": 0.704529,
     "end_time": "2023-10-08T08:43:51.053135",
     "exception": false,
     "start_time": "2023-10-08T08:43:50.348606",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sentence_pair_id</th>\n",
       "      <th>yjcaptions_id</th>\n",
       "      <th>sentence1</th>\n",
       "      <th>sentence2</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>100312_421853-104611-31624</td>\n",
       "      <td>レンガの建物の前を、乳母車を押した女性が歩いています。</td>\n",
       "      <td>厩舎で馬と女性とが寄り添っています。</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  sentence_pair_id               yjcaptions_id                    sentence1  \\\n",
       "0                0  100312_421853-104611-31624  レンガの建物の前を、乳母車を押した女性が歩いています。   \n",
       "\n",
       "            sentence2  label  \n",
       "0  厩舎で馬と女性とが寄り添っています。    0.0  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "import pandas as pd\n",
    "from urllib.request import urlopen\n",
    "\n",
    "df = pd.DataFrame([json.loads(line) for line in urlopen(jsts_url).readlines()])\n",
    "df.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a83a2374",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.230679Z",
     "start_time": "2023-10-07T08:50:54.218161Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:51.062681Z",
     "iopub.status.busy": "2023-10-08T08:43:51.061676Z",
     "iopub.status.idle": "2023-10-08T08:43:51.067684Z",
     "shell.execute_reply": "2023-10-08T08:43:51.067684Z"
    },
    "papermill": {
     "duration": 0.011549,
     "end_time": "2023-10-08T08:43:51.069685",
     "exception": false,
     "start_time": "2023-10-08T08:43:51.058136",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1457, 5)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f178b04",
   "metadata": {
    "papermill": {
     "duration": 0.003349,
     "end_time": "2023-10-08T08:43:51.076563",
     "exception": false,
     "start_time": "2023-10-08T08:43:51.073214",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c2253436",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.800419Z",
     "start_time": "2023-10-07T08:50:54.234207Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:51.083078Z",
     "iopub.status.busy": "2023-10-08T08:43:51.083078Z",
     "iopub.status.idle": "2023-10-08T08:43:54.598822Z",
     "shell.execute_reply": "2023-10-08T08:43:54.598822Z"
    },
    "papermill": {
     "duration": 3.521258,
     "end_time": "2023-10-08T08:43:54.600832",
     "exception": false,
     "start_time": "2023-10-08T08:43:51.079574",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((1457, 768), (1457, 768))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentence1_embs = model.encode(query_prefix + df[\"sentence1\"])\n",
    "sentence2_embs = model.encode(query_prefix + df[\"sentence2\"])\n",
    "sentence1_embs.shape, sentence2_embs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d14cd38",
   "metadata": {
    "papermill": {
     "duration": 0.002998,
     "end_time": "2023-10-08T08:43:54.607999",
     "exception": false,
     "start_time": "2023-10-08T08:43:54.605001",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "31c1f3b0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.847350Z",
     "start_time": "2023-10-07T08:50:57.801925Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:54.616524Z",
     "iopub.status.busy": "2023-10-08T08:43:54.615527Z",
     "iopub.status.idle": "2023-10-08T08:43:54.661817Z",
     "shell.execute_reply": "2023-10-08T08:43:54.661817Z"
    },
    "papermill": {
     "duration": 0.051319,
     "end_time": "2023-10-08T08:43:54.662825",
     "exception": false,
     "start_time": "2023-10-08T08:43:54.611506",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7657650438247267"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from scipy.spatial.distance import cosine, euclidean\n",
    "from scipy.stats import spearmanr\n",
    "\n",
    "df[\"similarity\"] = [\n",
    "    1 - cosine(s1, s2) for s1, s2 in zip(sentence1_embs, sentence2_embs)\n",
    "]\n",
    "jsts_score = spearmanr(df[\"similarity\"], df[\"label\"])[0]\n",
    "jsts_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "255586f8",
   "metadata": {
    "papermill": {
     "duration": 0.004,
     "end_time": "2023-10-08T08:43:54.670822",
     "exception": false,
     "start_time": "2023-10-08T08:43:54.666822",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSICK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d63e3ad0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.645301Z",
     "start_time": "2023-10-07T08:51:00.119204Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:54.678851Z",
     "iopub.status.busy": "2023-10-08T08:43:54.677848Z",
     "iopub.status.idle": "2023-10-08T08:43:55.504784Z",
     "shell.execute_reply": "2023-10-08T08:43:55.503863Z"
    },
    "papermill": {
     "duration": 0.831893,
     "end_time": "2023-10-08T08:43:55.505740",
     "exception": false,
     "start_time": "2023-10-08T08:43:54.673847",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>pair_ID</th>\n",
       "      <th>data</th>\n",
       "      <th>sentence_A_En</th>\n",
       "      <th>sentence_B_En</th>\n",
       "      <th>entailment_label_En</th>\n",
       "      <th>relatedness_score_En</th>\n",
       "      <th>corr_entailment_labelAB_En</th>\n",
       "      <th>corr_entailment_labelBA_En</th>\n",
       "      <th>sentence_A_Ja</th>\n",
       "      <th>sentence_B_Ja</th>\n",
       "      <th>entailment_label_Ja</th>\n",
       "      <th>relatedness_score_Ja</th>\n",
       "      <th>image_ID</th>\n",
       "      <th>original_caption</th>\n",
       "      <th>semtag_short</th>\n",
       "      <th>semtag_long</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>6</td>\n",
       "      <td>test</td>\n",
       "      <td>There is no boy playing outdoors and there is ...</td>\n",
       "      <td>A group of kids is playing in a yard and an ol...</td>\n",
       "      <td>neutral</td>\n",
       "      <td>3.3</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>戸外で遊んでいる男の子は一人もおらず、微笑んでいる男性は一人もいない</td>\n",
       "      <td>子供たちのグループが庭で遊んでいて、後ろの方には年を取った男性が立っている</td>\n",
       "      <td>contradiction</td>\n",
       "      <td>2.3</td>\n",
       "      <td>3155657768_b83a7831e5.jpg</td>\n",
       "      <td>The children are playing outdoors , while a ma...</td>\n",
       "      <td>Negation#Numerical</td>\n",
       "      <td>Numerical;人;名詞,接尾,助数詞,*#Negation;ない;助動詞,*,*,*#...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   pair_ID  data                                      sentence_A_En  \\\n",
       "0        6  test  There is no boy playing outdoors and there is ...   \n",
       "\n",
       "                                       sentence_B_En entailment_label_En  \\\n",
       "0  A group of kids is playing in a yard and an ol...             neutral   \n",
       "\n",
       "   relatedness_score_En corr_entailment_labelAB_En corr_entailment_labelBA_En  \\\n",
       "0                   3.3                        NaN                        NaN   \n",
       "\n",
       "                        sentence_A_Ja                          sentence_B_Ja  \\\n",
       "0  戸外で遊んでいる男の子は一人もおらず、微笑んでいる男性は一人もいない  子供たちのグループが庭で遊んでいて、後ろの方には年を取った男性が立っている   \n",
       "\n",
       "  entailment_label_Ja  relatedness_score_Ja                   image_ID  \\\n",
       "0       contradiction                   2.3  3155657768_b83a7831e5.jpg   \n",
       "\n",
       "                                    original_caption        semtag_short  \\\n",
       "0  The children are playing outdoors , while a ma...  Negation#Numerical   \n",
       "\n",
       "                                         semtag_long  \n",
       "0  Numerical;人;名詞,接尾,助数詞,*#Negation;ない;助動詞,*,*,*#...  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(jsick_url, sep=\"\\t\")\n",
    "df.head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "daf022a6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.660322Z",
     "start_time": "2023-10-07T08:51:00.647809Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:55.515783Z",
     "iopub.status.busy": "2023-10-08T08:43:55.515783Z",
     "iopub.status.idle": "2023-10-08T08:43:55.535876Z",
     "shell.execute_reply": "2023-10-08T08:43:55.534943Z"
    },
    "papermill": {
     "duration": 0.026137,
     "end_time": "2023-10-08T08:43:55.536876",
     "exception": false,
     "start_time": "2023-10-08T08:43:55.510739",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4927, 16)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2e89c25",
   "metadata": {
    "papermill": {
     "duration": 0.004559,
     "end_time": "2023-10-08T08:43:55.544968",
     "exception": false,
     "start_time": "2023-10-08T08:43:55.540409",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f249570f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:06.645910Z",
     "start_time": "2023-10-07T08:51:01.623993Z"
    },
    "execution": {
     "iopub.execute_input": "2023-10-08T08:43:55.552486Z",
     "iopub.status.busy": "2023-10-08T08:43:55.552486Z",
     "iopub.status.idle": "2023-10-08T08:44:00.331562Z",
     "shell.execute_reply": "2023-10-08T08:44:00.331562Z"
    },
    "papermill": {
     "duration": 4.785174,
     "end_time": "2023-10-08T08:44:00.333142",
     "exception": false,
     "start_time": "2023-10-08T08:43:55.547968",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((4927, 768), (4927, 768))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sentence1_embs = model.encode(query_prefix + df[\"sentence_A_Ja\"])\n",
    "sentence2_embs = model.encode(query_prefix + df[\"sentence_B_Ja\"])\n",
    "sentence1_embs.shape, sentence2_embs.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5a6d152",
   "metadata": {
    "papermill": {
     "duration": 0.003006,
     "end_time": "2023-10-08T08:44:00.340115",
     "exception": false,
     "start_time": "2023-10-08T08:44:00.337109",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "46b45fd5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:44:00.349123Z",
     "iopub.status.busy": "2023-10-08T08:44:00.349123Z",
     "iopub.status.idle": "2023-10-08T08:44:00.489822Z",
     "shell.execute_reply": "2023-10-08T08:44:00.488820Z"
    },
    "papermill": {
     "duration": 0.147217,
     "end_time": "2023-10-08T08:44:00.490855",
     "exception": false,
     "start_time": "2023-10-08T08:44:00.343638",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6524583232802058"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from scipy.spatial.distance import cosine\n",
    "from scipy.stats import spearmanr\n",
    "\n",
    "df[\"similarity\"] = [\n",
    "    1 - cosine(s1, s2) for s1, s2 in zip(sentence1_embs, sentence2_embs)\n",
    "]\n",
    "jsick_score = spearmanr(df[\"similarity\"], df[\"relatedness_score_Ja\"])[0]\n",
    "jsick_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c930aab3",
   "metadata": {
    "papermill": {
     "duration": 0.004006,
     "end_time": "2023-10-08T08:44:00.498979",
     "exception": false,
     "start_time": "2023-10-08T08:44:00.494973",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Miracle\n",
    "* Need access token for huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "4f8bdb96",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:44:00.508546Z",
     "iopub.status.busy": "2023-10-08T08:44:00.508546Z",
     "iopub.status.idle": "2023-10-08T08:44:00.519999Z",
     "shell.execute_reply": "2023-10-08T08:44:00.519999Z"
    },
    "papermill": {
     "duration": 0.018007,
     "end_time": "2023-10-08T08:44:00.521505",
     "exception": false,
     "start_time": "2023-10-08T08:44:00.503498",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import dotenv\n",
    "\n",
    "dotenv.load_dotenv(\"huggingface_access_token\", override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2d4aa8f7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:44:00.531008Z",
     "iopub.status.busy": "2023-10-08T08:44:00.529968Z",
     "iopub.status.idle": "2023-10-08T08:44:03.535411Z",
     "shell.execute_reply": "2023-10-08T08:44:03.534448Z"
    },
    "papermill": {
     "duration": 3.010898,
     "end_time": "2023-10-08T08:44:03.536412",
     "exception": false,
     "start_time": "2023-10-08T08:44:00.525514",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset miracl (G:/cache/miracl___miracl/ja/1.0.0/f598b4ee332f2b16e82c6c83ab1ba82e1a7777ef82e7ce3c1416f6b20a142313)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['query_id', 'query', 'positive_passages', 'negative_passages'],\n",
       "    num_rows: 860\n",
       "})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import datasets\n",
    "\n",
    "# query and positives\n",
    "ds = datasets.load_dataset(\n",
    "    \"miracl/miracl\", \"ja\", use_auth_token=os.environ[\"HF_ACCESS_TOKEN\"], split=\"dev\"\n",
    ")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a68d91b8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:44:03.546979Z",
     "iopub.status.busy": "2023-10-08T08:44:03.545979Z",
     "iopub.status.idle": "2023-10-08T08:44:05.674727Z",
     "shell.execute_reply": "2023-10-08T08:44:05.674727Z"
    },
    "papermill": {
     "duration": 2.136288,
     "end_time": "2023-10-08T08:44:05.676728",
     "exception": false,
     "start_time": "2023-10-08T08:44:03.540440",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset miracl-corpus (G:/cache/miracl___miracl-corpus/ja/1.0.0/16b566312c83a2e1f94d0813c8702b464b97f6b8959336adf062d289ce9b51fa)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7b0a311dbcdf4262912d5cf18050fe20",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['docid', 'title', 'text'],\n",
       "        num_rows: 6953614\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# all corpus texts\n",
    "corpus = datasets.load_dataset(\"miracl/miracl-corpus\", \"ja\")\n",
    "corpus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "33a9c4af",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:44:05.687250Z",
     "iopub.status.busy": "2023-10-08T08:44:05.687250Z",
     "iopub.status.idle": "2023-10-08T08:44:05.862471Z",
     "shell.execute_reply": "2023-10-08T08:44:05.862471Z"
    },
    "papermill": {
     "duration": 0.183748,
     "end_time": "2023-10-08T08:44:05.864476",
     "exception": false,
     "start_time": "2023-10-08T08:44:05.680728",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(860,\n",
       " ['0', '3', '4', '5', '7'],\n",
       " dict_keys(['docids', 'indices']),\n",
       " ['2681119#0', '2681119#1'],\n",
       " [1393435, 1393436])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hard negatives\n",
    "with open(\"./miracl_hard_negs_1000.json\") as f:\n",
    "    hn = json.loads(f.read())\n",
    "len(hn), list(hn.keys())[:5], hn[\"0\"].keys(), hn[\"0\"][\"docids\"][:2], hn[\"0\"][\"indices\"][\n",
    "    :2\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "825729ee",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T08:44:05.875991Z",
     "iopub.status.busy": "2023-10-08T08:44:05.875991Z",
     "iopub.status.idle": "2023-10-08T09:00:04.476876Z",
     "shell.execute_reply": "2023-10-08T09:00:04.475921Z"
    },
    "papermill": {
     "duration": 958.608434,
     "end_time": "2023-10-08T09:00:04.477912",
     "exception": false,
     "start_time": "2023-10-08T08:44:05.869478",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1790, 583, 0.3256983240223464)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.spatial.distance import cdist\n",
    "\n",
    "\n",
    "def get_text(corpus_item):\n",
    "    return corpus_item[\"title\"] + \" \" + corpus_item[\"text\"]\n",
    "\n",
    "\n",
    "n_total_pos = 0\n",
    "n_total_tp = 0\n",
    "\n",
    "for item in ds:\n",
    "    # query\n",
    "    query_emb = model.encode([query_prefix + item[\"query\"]])\n",
    "\n",
    "    # passages are set(300 hard negatives + positives)\n",
    "    positive_docids = [pp[\"docid\"] for pp in item[\"positive_passages\"]]\n",
    "    positive_texts = [get_text(pp) for pp in item[\"positive_passages\"]]\n",
    "    hn_docids = hn[item[\"query_id\"]][\"docids\"][:miracle_n_hard_negs]\n",
    "    hn_indices = hn[item[\"query_id\"]][\"indices\"][:miracle_n_hard_negs]\n",
    "\n",
    "    # drop hard negatives in positives\n",
    "    indices_to_keep = [\n",
    "        i for i, docid in enumerate(hn_docids) if docid not in positive_docids\n",
    "    ]\n",
    "    hn_docids = [hn_docids[i] for i in indices_to_keep]\n",
    "    hn_indices = [hn_indices[i] for i in indices_to_keep]\n",
    "\n",
    "    # search target\n",
    "    target_docids = positive_docids + hn_docids\n",
    "    target_texts = positive_texts + [\n",
    "        get_text(corpus[\"train\"][idx]) for idx in hn_indices\n",
    "    ]\n",
    "\n",
    "    assert set([corpus[\"train\"][idx][\"docid\"] for idx in hn_indices]) == set(hn_docids)\n",
    "\n",
    "    # embedding\n",
    "    target_embs = model.encode([passage_prefix + text for text in target_texts])\n",
    "\n",
    "    # topK\n",
    "    topk_indices = np.argsort(cdist(query_emb, target_embs, metric=\"cosine\"))[0][\n",
    "        :miracle_n_recall\n",
    "    ]\n",
    "\n",
    "    n_pos = len(positive_docids)\n",
    "    n_tp = len(\n",
    "        set(topk_indices) & set(range(len(positive_docids)))\n",
    "    )  # positives are first indices\n",
    "\n",
    "    n_total_pos += n_pos\n",
    "    n_total_tp += n_tp\n",
    "\n",
    "    # if n_pos > n_tp:\n",
    "    # print(f\"{item['query_id']}:{n_tp}/{n_pos}\", end=\", \")\n",
    "\n",
    "miracl_recall = n_total_tp / n_total_pos\n",
    "\n",
    "n_total_pos, n_total_tp, miracl_recall"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4213f2b1",
   "metadata": {
    "papermill": {
     "duration": 0.005005,
     "end_time": "2023-10-08T09:00:04.487420",
     "exception": false,
     "start_time": "2023-10-08T09:00:04.482415",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5d65b84c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T09:00:04.500950Z",
     "iopub.status.busy": "2023-10-08T09:00:04.500950Z",
     "iopub.status.idle": "2023-10-08T09:00:04.508489Z",
     "shell.execute_reply": "2023-10-08T09:00:04.507472Z"
    },
    "papermill": {
     "duration": 0.017535,
     "end_time": "2023-10-08T09:00:04.510474",
     "exception": false,
     "start_time": "2023-10-08T09:00:04.492939",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('MU-Kindai/SBERT-JSNLI-base',\n",
       " 0.7657650438247267,\n",
       " 0.6524583232802058,\n",
       " 0.3256983240223464)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_id, jsts_score, jsick_score, miracl_recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1ec4d88b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-10-08T09:00:04.522502Z",
     "iopub.status.busy": "2023-10-08T09:00:04.522502Z",
     "iopub.status.idle": "2023-10-08T09:00:04.541547Z",
     "shell.execute_reply": "2023-10-08T09:00:04.539544Z"
    },
    "papermill": {
     "duration": 0.028057,
     "end_time": "2023-10-08T09:00:04.543049",
     "exception": false,
     "start_time": "2023-10-08T09:00:04.514992",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(f'./scores/{model_id.replace(\"/\", \"_\")}.txt', \"w\") as f:\n",
    "    f.write(\n",
    "        json.dumps(\n",
    "            {\n",
    "                \"model_id\": model_id,\n",
    "                \"jsts\": jsts_score,\n",
    "                \"jsick\": jsick_score,\n",
    "                \"miracl\": miracl_recall,\n",
    "            }\n",
    "        )\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.11"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 980.117835,
   "end_time": "2023-10-08T09:00:05.743852",
   "environment_variables": {},
   "exception": null,
   "input_path": "template_sentence-transformers.ipynb",
   "output_path": "2_MU-Kindai_SBERT-JSNLI-base.ipynb",
   "parameters": {
    "model_id": "MU-Kindai/SBERT-JSNLI-base",
    "passage_prefix": "",
    "query_prefix": ""
   },
   "start_time": "2023-10-08T08:43:45.626017",
   "version": "2.4.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "02c2702a58ea2fe60fd5f26dd152a70e7993d77024040a4f035d0ea16923b730"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "2634f63ed790418eb4cf6f13f11d3fce": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "3b04534234d84e8cab1af236df7e70a8": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_e61cd77bcd2047caa6037abb693430e7",
       "placeholder": "​",
       "style": "IPY_MODEL_a68685a7f5d045979f017a9f2f7a4710",
       "tabbable": null,
       "tooltip": null,
       "value": " 1/1 [00:00&lt;00:00,  2.45it/s]"
      }
     },
     "3c6fd3b4fda043c6860543c34fc1386a": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "7b0a311dbcdf4262912d5cf18050fe20": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_853b7bb4dbc849588d0625c1dc7668c3",
        "IPY_MODEL_c4ac60f5b7614727b851e41ac46e0ec5",
        "IPY_MODEL_3b04534234d84e8cab1af236df7e70a8"
       ],
       "layout": "IPY_MODEL_b19566a09aac4510ac032203aa7bbd77",
       "tabbable": null,
       "tooltip": null
      }
     },
     "853b7bb4dbc849588d0625c1dc7668c3": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_2634f63ed790418eb4cf6f13f11d3fce",
       "placeholder": "​",
       "style": "IPY_MODEL_d807af0a278a41919b2de57628687046",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "a68685a7f5d045979f017a9f2f7a4710": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "b19566a09aac4510ac032203aa7bbd77": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "bc8651eb071b40d18384ca600080f785": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "c4ac60f5b7614727b851e41ac46e0ec5": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_3c6fd3b4fda043c6860543c34fc1386a",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_bc8651eb071b40d18384ca600080f785",
       "tabbable": null,
       "tooltip": null,
       "value": 1.0
      }
     },
     "d807af0a278a41919b2de57628687046": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "e61cd77bcd2047caa6037abb693430e7": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}