1_colorfulscoop_sbert-base-ja.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bbae78c5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.949996Z",
     "start_time": "2023-10-07T08:50:44.497041Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:46.914510Z",
     "iopub.status.busy": "2023-11-10T04:33:46.914510Z",
     "iopub.status.idle": "2023-11-10T04:33:49.564784Z",
     "shell.execute_reply": "2023-11-10T04:33:49.564784Z"
    },
    "papermill": {
     "duration": 2.656296,
     "end_time": "2023-11-10T04:33:49.566805",
     "exception": false,
     "start_time": "2023-11-10T04:33:46.910509",
     "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": "87076193",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.965421Z",
     "start_time": "2023-10-07T08:50:45.951998Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:49.574781Z",
     "iopub.status.busy": "2023-11-10T04:33:49.573780Z",
     "iopub.status.idle": "2023-11-10T04:33:49.581780Z",
     "shell.execute_reply": "2023-11-10T04:33:49.580785Z"
    },
    "papermill": {
     "duration": 0.012984,
     "end_time": "2023-11-10T04:33:49.582795",
     "exception": false,
     "start_time": "2023-11-10T04:33:49.569811",
     "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": "284711e1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:49.590782Z",
     "iopub.status.busy": "2023-11-10T04:33:49.589792Z",
     "iopub.status.idle": "2023-11-10T04:33:49.596826Z",
     "shell.execute_reply": "2023-11-10T04:33:49.596826Z"
    },
    "papermill": {
     "duration": 0.012039,
     "end_time": "2023-11-10T04:33:49.597820",
     "exception": false,
     "start_time": "2023-11-10T04:33:49.585781",
     "status": "completed"
    },
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "model_id = \"colorfulscoop/sbert-base-ja\"\n",
    "query_prefix = \"\"\n",
    "passage_prefix = \"\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06e64eac",
   "metadata": {
    "papermill": {
     "duration": 0.003481,
     "end_time": "2023-11-10T04:33:49.604306",
     "exception": false,
     "start_time": "2023-11-10T04:33:49.600825",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "61085709",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:53.766453Z",
     "start_time": "2023-10-07T08:50:51.631864Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:49.610313Z",
     "iopub.status.busy": "2023-11-10T04:33:49.610313Z",
     "iopub.status.idle": "2023-11-10T04:33:54.130734Z",
     "shell.execute_reply": "2023-11-10T04:33:54.129731Z"
    },
    "papermill": {
     "duration": 4.524423,
     "end_time": "2023-11-10T04:33:54.131735",
     "exception": false,
     "start_time": "2023-11-10T04:33:49.607312",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer(model_id)\n",
    "model.max_seq_length = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2de6b5ca",
   "metadata": {
    "papermill": {
     "duration": 0.004001,
     "end_time": "2023-11-10T04:33:54.138735",
     "exception": false,
     "start_time": "2023-11-10T04:33:54.134734",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "83cd7571",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.215161Z",
     "start_time": "2023-10-07T08:50:53.767452Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:54.145736Z",
     "iopub.status.busy": "2023-11-10T04:33:54.145736Z",
     "iopub.status.idle": "2023-11-10T04:33:55.598023Z",
     "shell.execute_reply": "2023-11-10T04:33:55.598023Z"
    },
    "papermill": {
     "duration": 1.458282,
     "end_time": "2023-11-10T04:33:55.600018",
     "exception": false,
     "start_time": "2023-11-10T04:33:54.141736",
     "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": "ced93339",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.230679Z",
     "start_time": "2023-10-07T08:50:54.218161Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:55.607976Z",
     "iopub.status.busy": "2023-11-10T04:33:55.607976Z",
     "iopub.status.idle": "2023-11-10T04:33:55.612988Z",
     "shell.execute_reply": "2023-11-10T04:33:55.612988Z"
    },
    "papermill": {
     "duration": 0.011025,
     "end_time": "2023-11-10T04:33:55.614013",
     "exception": false,
     "start_time": "2023-11-10T04:33:55.602988",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1457, 5)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b9521d4",
   "metadata": {
    "papermill": {
     "duration": 0.003991,
     "end_time": "2023-11-10T04:33:55.621013",
     "exception": false,
     "start_time": "2023-11-10T04:33:55.617022",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "68b680fe",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.800419Z",
     "start_time": "2023-10-07T08:50:54.234207Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:55.628014Z",
     "iopub.status.busy": "2023-11-10T04:33:55.628014Z",
     "iopub.status.idle": "2023-11-10T04:33:58.875201Z",
     "shell.execute_reply": "2023-11-10T04:33:58.875201Z"
    },
    "papermill": {
     "duration": 3.253165,
     "end_time": "2023-11-10T04:33:58.877185",
     "exception": false,
     "start_time": "2023-11-10T04:33:55.624020",
     "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": "f20b587e",
   "metadata": {
    "papermill": {
     "duration": 0.002993,
     "end_time": "2023-11-10T04:33:58.884177",
     "exception": false,
     "start_time": "2023-11-10T04:33:58.881184",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ae5c1d7f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.847350Z",
     "start_time": "2023-10-07T08:50:57.801925Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:58.893169Z",
     "iopub.status.busy": "2023-11-10T04:33:58.893169Z",
     "iopub.status.idle": "2023-11-10T04:33:58.939201Z",
     "shell.execute_reply": "2023-11-10T04:33:58.939201Z"
    },
    "papermill": {
     "duration": 0.051079,
     "end_time": "2023-11-10T04:33:58.940249",
     "exception": false,
     "start_time": "2023-11-10T04:33:58.889170",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7423951868203088"
      ]
     },
     "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": "66b08a0d",
   "metadata": {
    "papermill": {
     "duration": 0.00326,
     "end_time": "2023-11-10T04:33:58.947766",
     "exception": false,
     "start_time": "2023-11-10T04:33:58.944506",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSICK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "43ca66fa",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.645301Z",
     "start_time": "2023-10-07T08:51:00.119204Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:58.956773Z",
     "iopub.status.busy": "2023-11-10T04:33:58.955782Z",
     "iopub.status.idle": "2023-11-10T04:33:59.940030Z",
     "shell.execute_reply": "2023-11-10T04:33:59.940030Z"
    },
    "papermill": {
     "duration": 0.991256,
     "end_time": "2023-11-10T04:33:59.942028",
     "exception": false,
     "start_time": "2023-11-10T04:33:58.950772",
     "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": "65f26905",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.660322Z",
     "start_time": "2023-10-07T08:51:00.647809Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:59.950547Z",
     "iopub.status.busy": "2023-11-10T04:33:59.950547Z",
     "iopub.status.idle": "2023-11-10T04:33:59.955553Z",
     "shell.execute_reply": "2023-11-10T04:33:59.955553Z"
    },
    "papermill": {
     "duration": 0.01102,
     "end_time": "2023-11-10T04:33:59.956559",
     "exception": false,
     "start_time": "2023-11-10T04:33:59.945539",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4927, 16)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d94efe2",
   "metadata": {
    "papermill": {
     "duration": 0.002999,
     "end_time": "2023-11-10T04:33:59.963545",
     "exception": false,
     "start_time": "2023-11-10T04:33:59.960546",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9071d59a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:06.645910Z",
     "start_time": "2023-10-07T08:51:01.623993Z"
    },
    "execution": {
     "iopub.execute_input": "2023-11-10T04:33:59.972550Z",
     "iopub.status.busy": "2023-11-10T04:33:59.971547Z",
     "iopub.status.idle": "2023-11-10T04:34:03.908610Z",
     "shell.execute_reply": "2023-11-10T04:34:03.908610Z"
    },
    "papermill": {
     "duration": 3.943102,
     "end_time": "2023-11-10T04:34:03.910648",
     "exception": false,
     "start_time": "2023-11-10T04:33:59.967546",
     "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": "0b342b11",
   "metadata": {
    "papermill": {
     "duration": 0.004002,
     "end_time": "2023-11-10T04:34:03.918620",
     "exception": false,
     "start_time": "2023-11-10T04:34:03.914618",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "00dcc9cf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:34:03.927618Z",
     "iopub.status.busy": "2023-11-10T04:34:03.927618Z",
     "iopub.status.idle": "2023-11-10T04:34:04.051203Z",
     "shell.execute_reply": "2023-11-10T04:34:04.051203Z"
    },
    "papermill": {
     "duration": 0.13059,
     "end_time": "2023-11-10T04:34:04.053215",
     "exception": false,
     "start_time": "2023-11-10T04:34:03.922625",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6573168898177302"
      ]
     },
     "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": "00cb1f78",
   "metadata": {
    "papermill": {
     "duration": 0.004992,
     "end_time": "2023-11-10T04:34:04.062199",
     "exception": false,
     "start_time": "2023-11-10T04:34:04.057207",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Miracle\n",
    "* Need access token for huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "99a0d1e6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:34:04.072195Z",
     "iopub.status.busy": "2023-11-10T04:34:04.071177Z",
     "iopub.status.idle": "2023-11-10T04:34:04.098247Z",
     "shell.execute_reply": "2023-11-10T04:34:04.098247Z"
    },
    "papermill": {
     "duration": 0.034078,
     "end_time": "2023-11-10T04:34:04.100255",
     "exception": false,
     "start_time": "2023-11-10T04:34:04.066177",
     "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": "12388582",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:34:04.108654Z",
     "iopub.status.busy": "2023-11-10T04:34:04.108654Z",
     "iopub.status.idle": "2023-11-10T04:34:06.678673Z",
     "shell.execute_reply": "2023-11-10T04:34:06.677701Z"
    },
    "papermill": {
     "duration": 2.575428,
     "end_time": "2023-11-10T04:34:06.679676",
     "exception": false,
     "start_time": "2023-11-10T04:34:04.104248",
     "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": "9bd121ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:34:06.689687Z",
     "iopub.status.busy": "2023-11-10T04:34:06.689687Z",
     "iopub.status.idle": "2023-11-10T04:34:09.002713Z",
     "shell.execute_reply": "2023-11-10T04:34:09.001707Z"
    },
    "papermill": {
     "duration": 2.319031,
     "end_time": "2023-11-10T04:34:09.003712",
     "exception": false,
     "start_time": "2023-11-10T04:34:06.684681",
     "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": "861d130ef8f34558babd7dcb1259bf9a",
       "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": "d6d7f9e9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:34:09.013709Z",
     "iopub.status.busy": "2023-11-10T04:34:09.012708Z",
     "iopub.status.idle": "2023-11-10T04:34:09.177710Z",
     "shell.execute_reply": "2023-11-10T04:34:09.176709Z"
    },
    "papermill": {
     "duration": 0.171,
     "end_time": "2023-11-10T04:34:09.179709",
     "exception": false,
     "start_time": "2023-11-10T04:34:09.008709",
     "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": "455ba307",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:34:09.191708Z",
     "iopub.status.busy": "2023-11-10T04:34:09.190710Z",
     "iopub.status.idle": "2023-11-10T04:48:34.195690Z",
     "shell.execute_reply": "2023-11-10T04:48:34.195690Z"
    },
    "papermill": {
     "duration": 865.012026,
     "end_time": "2023-11-10T04:48:34.197738",
     "exception": false,
     "start_time": "2023-11-10T04:34:09.185712",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1790, 455, 0.2541899441340782)"
      ]
     },
     "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": "6ba8bd0c",
   "metadata": {
    "papermill": {
     "duration": 0.004999,
     "end_time": "2023-11-10T04:48:34.206695",
     "exception": false,
     "start_time": "2023-11-10T04:48:34.201696",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "22ef6c7f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:48:34.216697Z",
     "iopub.status.busy": "2023-11-10T04:48:34.216697Z",
     "iopub.status.idle": "2023-11-10T04:48:34.226703Z",
     "shell.execute_reply": "2023-11-10T04:48:34.226703Z"
    },
    "papermill": {
     "duration": 0.016999,
     "end_time": "2023-11-10T04:48:34.228731",
     "exception": false,
     "start_time": "2023-11-10T04:48:34.211732",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('colorfulscoop/sbert-base-ja',\n",
       " 0.7423951868203088,\n",
       " 0.6573168898177302,\n",
       " 0.2541899441340782)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_id, jsts_score, jsick_score, miracl_recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4ed82eaf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-10T04:48:34.237695Z",
     "iopub.status.busy": "2023-11-10T04:48:34.237695Z",
     "iopub.status.idle": "2023-11-10T04:48:34.243697Z",
     "shell.execute_reply": "2023-11-10T04:48:34.242734Z"
    },
    "papermill": {
     "duration": 0.012023,
     "end_time": "2023-11-10T04:48:34.244729",
     "exception": false,
     "start_time": "2023-11-10T04:48:34.232706",
     "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": 891.365295,
   "end_time": "2023-11-10T04:48:36.456403",
   "environment_variables": {},
   "exception": null,
   "input_path": "template_sentence-transformers.ipynb",
   "output_path": "1_colorfulscoop_sbert-base-ja.ipynb",
   "parameters": {
    "model_id": "colorfulscoop/sbert-base-ja",
    "passage_prefix": "",
    "query_prefix": ""
   },
   "start_time": "2023-11-10T04:33:45.091108",
   "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": {
     "08322ac222e6493bbaae6d2be409dbbd": {
      "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_ec05e7f3cff44a6dbaa715b814238bcc",
       "max": 1.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_f4e37d7c202842bcb2f5d821d4a845b4",
       "tabbable": null,
       "tooltip": null,
       "value": 1.0
      }
     },
     "0c5f6df28f6f406e8a42e4100f02a1f8": {
      "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
      }
     },
     "31d1717974404b5589752e7bdfc878f8": {
      "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
      }
     },
     "861d130ef8f34558babd7dcb1259bf9a": {
      "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_8c71a6dc0a034e709f2c64ead5e2a201",
        "IPY_MODEL_08322ac222e6493bbaae6d2be409dbbd",
        "IPY_MODEL_bb0f7bc707e14a93a2ffcba57b5e38b4"
       ],
       "layout": "IPY_MODEL_997935bd839445ab8ce215769bb95130",
       "tabbable": null,
       "tooltip": null
      }
     },
     "8c71a6dc0a034e709f2c64ead5e2a201": {
      "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_da60653ddd734225ac7d01fdec5a7234",
       "placeholder": "​",
       "style": "IPY_MODEL_31d1717974404b5589752e7bdfc878f8",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "997935bd839445ab8ce215769bb95130": {
      "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
      }
     },
     "bb0f7bc707e14a93a2ffcba57b5e38b4": {
      "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_0c5f6df28f6f406e8a42e4100f02a1f8",
       "placeholder": "​",
       "style": "IPY_MODEL_cab2a944343b462aaa1648ed97149dca",
       "tabbable": null,
       "tooltip": null,
       "value": " 1/1 [00:00&lt;00:00,  2.30it/s]"
      }
     },
     "cab2a944343b462aaa1648ed97149dca": {
      "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
      }
     },
     "da60653ddd734225ac7d01fdec5a7234": {
      "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
      }
     },
     "ec05e7f3cff44a6dbaa715b814238bcc": {
      "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
      }
     },
     "f4e37d7c202842bcb2f5d821d4a845b4": {
      "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": ""
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}