3_MU-Kindai_SBERT-JSNLI-large.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "de257982",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.949996Z",
     "start_time": "2023-10-07T08:50:44.497041Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:32.306889Z",
     "iopub.status.busy": "2023-11-10T05:06:32.306889Z",
     "iopub.status.idle": "2023-11-10T05:06:33.596799Z",
     "shell.execute_reply": "2023-11-10T05:06:33.596799Z"
    },
    "papermill": {
     "duration": 1.29582,
     "end_time": "2023-11-10T05:06:33.598810",
     "exception": false,
     "start_time": "2023-11-10T05:06:32.302990",
     "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": "2868d72d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.965421Z",
     "start_time": "2023-10-07T08:50:45.951998Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:33.606797Z",
     "iopub.status.busy": "2023-11-10T05:06:33.605797Z",
     "iopub.status.idle": "2023-11-10T05:06:33.612854Z",
     "shell.execute_reply": "2023-11-10T05:06:33.612854Z"
    },
    "papermill": {
     "duration": 0.013103,
     "end_time": "2023-11-10T05:06:33.614899",
     "exception": false,
     "start_time": "2023-11-10T05:06:33.601796",
     "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": "d1fca540",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:33.620905Z",
     "iopub.status.busy": "2023-11-10T05:06:33.620905Z",
     "iopub.status.idle": "2023-11-10T05:06:33.629003Z",
     "shell.execute_reply": "2023-11-10T05:06:33.628443Z"
    },
    "papermill": {
     "duration": 0.01371,
     "end_time": "2023-11-10T05:06:33.630615",
     "exception": false,
     "start_time": "2023-11-10T05:06:33.616905",
     "status": "completed"
    },
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "model_id = \"MU-Kindai/SBERT-JSNLI-large\"\n",
    "query_prefix = \"\"\n",
    "passage_prefix = \"\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a999246e",
   "metadata": {
    "papermill": {
     "duration": 0.003303,
     "end_time": "2023-11-10T05:06:33.636564",
     "exception": false,
     "start_time": "2023-11-10T05:06:33.633261",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1998a2a9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:53.766453Z",
     "start_time": "2023-10-07T08:50:51.631864Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:33.643571Z",
     "iopub.status.busy": "2023-11-10T05:06:33.643571Z",
     "iopub.status.idle": "2023-11-10T05:06:37.729178Z",
     "shell.execute_reply": "2023-11-10T05:06:37.728178Z"
    },
    "papermill": {
     "duration": 4.09161,
     "end_time": "2023-11-10T05:06:37.731178",
     "exception": false,
     "start_time": "2023-11-10T05:06:33.639568",
     "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": "861ff064",
   "metadata": {
    "papermill": {
     "duration": 0.004999,
     "end_time": "2023-11-10T05:06:37.739178",
     "exception": false,
     "start_time": "2023-11-10T05:06:37.734179",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9fc82bad",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.215161Z",
     "start_time": "2023-10-07T08:50:53.767452Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:37.748179Z",
     "iopub.status.busy": "2023-11-10T05:06:37.748179Z",
     "iopub.status.idle": "2023-11-10T05:06:38.378338Z",
     "shell.execute_reply": "2023-11-10T05:06:38.378338Z"
    },
    "papermill": {
     "duration": 0.637147,
     "end_time": "2023-11-10T05:06:38.380326",
     "exception": false,
     "start_time": "2023-11-10T05:06:37.743179",
     "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": "592633ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.230679Z",
     "start_time": "2023-10-07T08:50:54.218161Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:38.389329Z",
     "iopub.status.busy": "2023-11-10T05:06:38.388329Z",
     "iopub.status.idle": "2023-11-10T05:06:38.394339Z",
     "shell.execute_reply": "2023-11-10T05:06:38.394339Z"
    },
    "papermill": {
     "duration": 0.01199,
     "end_time": "2023-11-10T05:06:38.396327",
     "exception": false,
     "start_time": "2023-11-10T05:06:38.384337",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1457, 5)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "323a6d9b",
   "metadata": {
    "papermill": {
     "duration": 0.003991,
     "end_time": "2023-11-10T05:06:38.403327",
     "exception": false,
     "start_time": "2023-11-10T05:06:38.399336",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "43d1fc5e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.800419Z",
     "start_time": "2023-10-07T08:50:54.234207Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:38.411345Z",
     "iopub.status.busy": "2023-11-10T05:06:38.411345Z",
     "iopub.status.idle": "2023-11-10T05:06:45.388019Z",
     "shell.execute_reply": "2023-11-10T05:06:45.388019Z"
    },
    "papermill": {
     "duration": 6.983694,
     "end_time": "2023-11-10T05:06:45.390019",
     "exception": false,
     "start_time": "2023-11-10T05:06:38.406325",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((1457, 1024), (1457, 1024))"
      ]
     },
     "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": "9232f7e6",
   "metadata": {
    "papermill": {
     "duration": 0.004002,
     "end_time": "2023-11-10T05:06:45.398014",
     "exception": false,
     "start_time": "2023-11-10T05:06:45.394012",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c1bd744e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.847350Z",
     "start_time": "2023-10-07T08:50:57.801925Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:45.407059Z",
     "iopub.status.busy": "2023-11-10T05:06:45.406014Z",
     "iopub.status.idle": "2023-11-10T05:06:45.451923Z",
     "shell.execute_reply": "2023-11-10T05:06:45.451366Z"
    },
    "papermill": {
     "duration": 0.050418,
     "end_time": "2023-11-10T05:06:45.452429",
     "exception": false,
     "start_time": "2023-11-10T05:06:45.402011",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7740053633973121"
      ]
     },
     "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": "e46f67c7",
   "metadata": {
    "papermill": {
     "duration": 0.003002,
     "end_time": "2023-11-10T05:06:45.459436",
     "exception": false,
     "start_time": "2023-11-10T05:06:45.456434",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSICK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d35e13db",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.645301Z",
     "start_time": "2023-10-07T08:51:00.119204Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:45.467438Z",
     "iopub.status.busy": "2023-11-10T05:06:45.467438Z",
     "iopub.status.idle": "2023-11-10T05:06:46.291111Z",
     "shell.execute_reply": "2023-11-10T05:06:46.290226Z"
    },
    "papermill": {
     "duration": 0.8287,
     "end_time": "2023-11-10T05:06:46.292136",
     "exception": false,
     "start_time": "2023-11-10T05:06:45.463436",
     "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": "4fc19541",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.660322Z",
     "start_time": "2023-10-07T08:51:00.647809Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:46.301104Z",
     "iopub.status.busy": "2023-11-10T05:06:46.301104Z",
     "iopub.status.idle": "2023-11-10T05:06:46.306110Z",
     "shell.execute_reply": "2023-11-10T05:06:46.306110Z"
    },
    "papermill": {
     "duration": 0.012,
     "end_time": "2023-11-10T05:06:46.308102",
     "exception": false,
     "start_time": "2023-11-10T05:06:46.296102",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4927, 16)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bf81c74",
   "metadata": {
    "papermill": {
     "duration": 0.004001,
     "end_time": "2023-11-10T05:06:46.315104",
     "exception": false,
     "start_time": "2023-11-10T05:06:46.311103",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1c59bee1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:06.645910Z",
     "start_time": "2023-10-07T08:51:01.623993Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:46.323103Z",
     "iopub.status.busy": "2023-11-10T05:06:46.323103Z",
     "iopub.status.idle": "2023-11-10T05:06:57.748440Z",
     "shell.execute_reply": "2023-11-10T05:06:57.747569Z"
    },
    "papermill": {
     "duration": 11.431357,
     "end_time": "2023-11-10T05:06:57.749467",
     "exception": false,
     "start_time": "2023-11-10T05:06:46.318110",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((4927, 1024), (4927, 1024))"
      ]
     },
     "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": "b5738541",
   "metadata": {
    "papermill": {
     "duration": 0.00399,
     "end_time": "2023-11-10T05:06:57.757440",
     "exception": false,
     "start_time": "2023-11-10T05:06:57.753450",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2eefe642",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:57.766485Z",
     "iopub.status.busy": "2023-11-10T05:06:57.766485Z",
     "iopub.status.idle": "2023-11-10T05:06:57.890605Z",
     "shell.execute_reply": "2023-11-10T05:06:57.890605Z"
    },
    "papermill": {
     "duration": 0.131175,
     "end_time": "2023-11-10T05:06:57.892613",
     "exception": false,
     "start_time": "2023-11-10T05:06:57.761438",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6768182045984232"
      ]
     },
     "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": "4722815b",
   "metadata": {
    "papermill": {
     "duration": 0.004019,
     "end_time": "2023-11-10T05:06:57.900642",
     "exception": false,
     "start_time": "2023-11-10T05:06:57.896623",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Miracle\n",
    "* Need access token for huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5c101433",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:57.909635Z",
     "iopub.status.busy": "2023-11-10T05:06:57.909635Z",
     "iopub.status.idle": "2023-11-10T05:06:57.921816Z",
     "shell.execute_reply": "2023-11-10T05:06:57.921816Z"
    },
    "papermill": {
     "duration": 0.018155,
     "end_time": "2023-11-10T05:06:57.922788",
     "exception": false,
     "start_time": "2023-11-10T05:06:57.904633",
     "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": "b5164412",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:57.931974Z",
     "iopub.status.busy": "2023-11-10T05:06:57.931974Z",
     "iopub.status.idle": "2023-11-10T05:06:59.852048Z",
     "shell.execute_reply": "2023-11-10T05:06:59.851039Z"
    },
    "papermill": {
     "duration": 1.926278,
     "end_time": "2023-11-10T05:06:59.853079",
     "exception": false,
     "start_time": "2023-11-10T05:06:57.926801",
     "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": "8976889e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:06:59.863056Z",
     "iopub.status.busy": "2023-11-10T05:06:59.863056Z",
     "iopub.status.idle": "2023-11-10T05:07:01.803353Z",
     "shell.execute_reply": "2023-11-10T05:07:01.802385Z"
    },
    "papermill": {
     "duration": 1.947328,
     "end_time": "2023-11-10T05:07:01.804390",
     "exception": false,
     "start_time": "2023-11-10T05:06:59.857062",
     "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": "0fc21298735040f584ae0724af0d26f2",
       "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": "1c47236d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:07:01.815352Z",
     "iopub.status.busy": "2023-11-10T05:07:01.814353Z",
     "iopub.status.idle": "2023-11-10T05:07:02.036779Z",
     "shell.execute_reply": "2023-11-10T05:07:02.036779Z"
    },
    "papermill": {
     "duration": 0.23035,
     "end_time": "2023-11-10T05:07:02.038702",
     "exception": false,
     "start_time": "2023-11-10T05:07:01.808352",
     "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": "139e99bd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:07:02.048691Z",
     "iopub.status.busy": "2023-11-10T05:07:02.048691Z",
     "iopub.status.idle": "2023-11-10T05:56:19.864554Z",
     "shell.execute_reply": "2023-11-10T05:56:19.863557Z"
    },
    "papermill": {
     "duration": 2957.844897,
     "end_time": "2023-11-10T05:56:19.887555",
     "exception": false,
     "start_time": "2023-11-10T05:07:02.042658",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1790, 498, 0.2782122905027933)"
      ]
     },
     "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",
    "corpus_dict = {item[\"docid\"]: get_text(item) for item in corpus[\"train\"]}\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",
    "\n",
    "    # drop hard negatives in positives\n",
    "    hn_docids = [docid for docid in hn_docids if docid not in positive_docids]\n",
    "\n",
    "    # search target\n",
    "    target_docids = positive_docids + hn_docids\n",
    "    target_texts = positive_texts + [corpus_dict[docid] for docid in 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": "eeede7bf",
   "metadata": {
    "papermill": {
     "duration": 0.006998,
     "end_time": "2023-11-10T05:56:19.905556",
     "exception": false,
     "start_time": "2023-11-10T05:56:19.898558",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5e1b2373",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:56:19.919555Z",
     "iopub.status.busy": "2023-11-10T05:56:19.919555Z",
     "iopub.status.idle": "2023-11-10T05:56:19.928559Z",
     "shell.execute_reply": "2023-11-10T05:56:19.927555Z"
    },
    "papermill": {
     "duration": 0.018,
     "end_time": "2023-11-10T05:56:19.929556",
     "exception": false,
     "start_time": "2023-11-10T05:56:19.911556",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('MU-Kindai/SBERT-JSNLI-large',\n",
       " 0.7740053633973121,\n",
       " 0.6768182045984232,\n",
       " 0.2782122905027933)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_id, jsts_score, jsick_score, miracl_recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d62a2381",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T05:56:19.941556Z",
     "iopub.status.busy": "2023-11-10T05:56:19.940556Z",
     "iopub.status.idle": "2023-11-10T05:56:19.960555Z",
     "shell.execute_reply": "2023-11-10T05:56:19.959557Z"
    },
    "papermill": {
     "duration": 0.026998,
     "end_time": "2023-11-10T05:56:19.961554",
     "exception": false,
     "start_time": "2023-11-10T05:56:19.934556",
     "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": 2991.586666,
   "end_time": "2023-11-10T05:56:22.598561",
   "environment_variables": {},
   "exception": null,
   "input_path": "template_sentence-transformers.ipynb",
   "output_path": "3_MU-Kindai_SBERT-JSNLI-large.ipynb",
   "parameters": {
    "model_id": "MU-Kindai/SBERT-JSNLI-large",
    "passage_prefix": "",
    "query_prefix": ""
   },
   "start_time": "2023-11-10T05:06:31.011895",
   "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": {
     "093f7d3eca604e62893039fa3be79fbb": {
      "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
      }
     },
     "0fc21298735040f584ae0724af0d26f2": {
      "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_ad0d68d3351842b49260c11d1c41771e",
        "IPY_MODEL_8649bbe1814c4e749e6bf2ea460f96a2",
        "IPY_MODEL_318f39fb35cc4155afc35b1acacb0425"
       ],
       "layout": "IPY_MODEL_093f7d3eca604e62893039fa3be79fbb",
       "tabbable": null,
       "tooltip": null
      }
     },
     "24ae3af535944ce7a75aeb5c982c599f": {
      "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
      }
     },
     "30deb3a9a5b541c584b85620f2c34b05": {
      "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
      }
     },
     "318f39fb35cc4155afc35b1acacb0425": {
      "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_24ae3af535944ce7a75aeb5c982c599f",
       "placeholder": "​",
       "style": "IPY_MODEL_30deb3a9a5b541c584b85620f2c34b05",
       "tabbable": null,
       "tooltip": null,
       "value": " 1/1 [00:00&lt;00:00,  2.61it/s]"
      }
     },
     "8649bbe1814c4e749e6bf2ea460f96a2": {
      "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_e47fbbcd396148c692aae3d2dfd40f34",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_d2c5aba3bfb64316874156e64bb8958c",
       "tabbable": null,
       "tooltip": null,
       "value": 1.0
      }
     },
     "a08f9a4459894a51b98f79242255147b": {
      "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
      }
     },
     "ad0d68d3351842b49260c11d1c41771e": {
      "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_a08f9a4459894a51b98f79242255147b",
       "placeholder": "​",
       "style": "IPY_MODEL_f87bf28846b2482392f9ebf333c51e88",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "d2c5aba3bfb64316874156e64bb8958c": {
      "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": ""
      }
     },
     "e47fbbcd396148c692aae3d2dfd40f34": {
      "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
      }
     },
     "f87bf28846b2482392f9ebf333c51e88": {
      "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
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}