34_pkshatech_GLuCoSE-base-ja-v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "edfc0f31",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.949996Z",
     "start_time": "2023-10-07T08:50:44.497041Z"
    },
    "papermill": {
     "duration": 1.328372,
     "end_time": "2023-11-10T10:09:03.422531",
     "exception": false,
     "start_time": "2023-11-10T10:09:02.094159",
     "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": "e4e27fb5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.965421Z",
     "start_time": "2023-10-07T08:50:45.951998Z"
    },
    "papermill": {
     "duration": 0.011817,
     "end_time": "2023-11-10T10:09:03.438357",
     "exception": false,
     "start_time": "2023-11-10T10:09:03.426540",
     "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": "11c3aa8a",
   "metadata": {
    "papermill": {
     "duration": 0.010638,
     "end_time": "2023-11-10T10:09:03.452466",
     "exception": false,
     "start_time": "2023-11-10T10:09:03.441828",
     "status": "completed"
    },
    "tags": [
     "injected-parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "model_id = \"pkshatech/GLuCoSE-base-ja-v2\"\n",
    "query_prefix = \"query: \"\n",
    "passage_prefix = \"passage: \"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84e7998f",
   "metadata": {
    "papermill": {
     "duration": 0.003,
     "end_time": "2023-11-10T10:09:03.458429",
     "exception": false,
     "start_time": "2023-11-10T10:09:03.455429",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f672a5e3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:53.766453Z",
     "start_time": "2023-10-07T08:50:51.631864Z"
    },
    "papermill": {
     "duration": 2.849968,
     "end_time": "2023-11-10T10:09:06.311439",
     "exception": false,
     "start_time": "2023-11-10T10:09:03.461471",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/oshizo/repo/venv/lib/python3.11/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:13: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm, trange\n",
      "/home/oshizo/repo/venv/lib/python3.11/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/home/oshizo/repo/venv/lib/python3.11/site-packages/torchvision/image.so: undefined symbol: _ZN3c1017RegisterOperatorsD1Ev'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?\n",
      "  warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f876114e4f914104ae4946252f4febd1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/9.96k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "model = SentenceTransformer(model_id, trust_remote_code=True)\n",
    "model.max_seq_length = 512"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12a92348",
   "metadata": {
    "papermill": {
     "duration": 0.003002,
     "end_time": "2023-11-10T10:09:06.318673",
     "exception": false,
     "start_time": "2023-11-10T10:09:06.315671",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "838c14db",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.215161Z",
     "start_time": "2023-10-07T08:50:53.767452Z"
    },
    "papermill": {
     "duration": 0.893847,
     "end_time": "2023-11-10T10:09:07.217033",
     "exception": false,
     "start_time": "2023-11-10T10:09:06.323186",
     "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": "e072e072",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:54.230679Z",
     "start_time": "2023-10-07T08:50:54.218161Z"
    },
    "papermill": {
     "duration": 0.012461,
     "end_time": "2023-11-10T10:09:07.232482",
     "exception": false,
     "start_time": "2023-11-10T10:09:07.220021",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1457, 5)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405ae7cb",
   "metadata": {
    "papermill": {
     "duration": 0.003521,
     "end_time": "2023-11-10T10:09:07.239015",
     "exception": false,
     "start_time": "2023-11-10T10:09:07.235494",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2570344b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.800419Z",
     "start_time": "2023-10-07T08:50:54.234207Z"
    },
    "papermill": {
     "duration": 3.158118,
     "end_time": "2023-11-10T10:09:10.400668",
     "exception": false,
     "start_time": "2023-11-10T10:09:07.242550",
     "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": "9730ed04",
   "metadata": {
    "papermill": {
     "duration": 0.003958,
     "end_time": "2023-11-10T10:09:10.409234",
     "exception": false,
     "start_time": "2023-11-10T10:09:10.405276",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f0a3606d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:57.847350Z",
     "start_time": "2023-10-07T08:50:57.801925Z"
    },
    "papermill": {
     "duration": 0.052785,
     "end_time": "2023-11-10T10:09:10.464524",
     "exception": false,
     "start_time": "2023-11-10T10:09:10.411739",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.808965162015302"
      ]
     },
     "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": "4a22f620",
   "metadata": {
    "papermill": {
     "duration": 0.004667,
     "end_time": "2023-11-10T10:09:10.473190",
     "exception": false,
     "start_time": "2023-11-10T10:09:10.468523",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# JSICK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eca38121",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.645301Z",
     "start_time": "2023-10-07T08:51:00.119204Z"
    },
    "papermill": {
     "duration": 0.889405,
     "end_time": "2023-11-10T10:09:11.365910",
     "exception": false,
     "start_time": "2023-11-10T10:09:10.476505",
     "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": "485fceb7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:00.660322Z",
     "start_time": "2023-10-07T08:51:00.647809Z"
    },
    "papermill": {
     "duration": 0.011534,
     "end_time": "2023-11-10T10:09:11.381422",
     "exception": false,
     "start_time": "2023-11-10T10:09:11.369888",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4927, 16)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be3c2725",
   "metadata": {
    "papermill": {
     "duration": 0.003997,
     "end_time": "2023-11-10T10:09:11.388524",
     "exception": false,
     "start_time": "2023-11-10T10:09:11.384527",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "64c9750d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:51:06.645910Z",
     "start_time": "2023-10-07T08:51:01.623993Z"
    },
    "papermill": {
     "duration": 3.911541,
     "end_time": "2023-11-10T10:09:15.304808",
     "exception": false,
     "start_time": "2023-11-10T10:09:11.393267",
     "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": "8a1dbded",
   "metadata": {
    "papermill": {
     "duration": 0.033047,
     "end_time": "2023-11-10T10:09:15.341849",
     "exception": false,
     "start_time": "2023-11-10T10:09:15.308802",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Correlation Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "79cc78b9",
   "metadata": {
    "papermill": {
     "duration": 0.145114,
     "end_time": "2023-11-10T10:09:15.491499",
     "exception": false,
     "start_time": "2023-11-10T10:09:15.346385",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8490230332181042"
      ]
     },
     "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": "2ac0e75d",
   "metadata": {
    "papermill": {
     "duration": 0.003999,
     "end_time": "2023-11-10T10:09:15.500419",
     "exception": false,
     "start_time": "2023-11-10T10:09:15.496420",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Miracle\n",
    "* Need access token for huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9f00c7a7",
   "metadata": {
    "papermill": {
     "duration": 0.019154,
     "end_time": "2023-11-10T10:09:15.523633",
     "exception": false,
     "start_time": "2023-11-10T10:09:15.504479",
     "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": "8d0dbf84",
   "metadata": {
    "papermill": {
     "duration": 2.222354,
     "end_time": "2023-11-10T10:09:17.748989",
     "exception": false,
     "start_time": "2023-11-10T10:09:15.526635",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/oshizo/repo/venv/lib/python3.11/site-packages/datasets/load.py:2554: FutureWarning: 'use_auth_token' was deprecated in favor of 'token' in version 2.14.0 and will be removed in 3.0.0.\n",
      "You can remove this warning by passing 'token=<use_auth_token>' instead.\n",
      "  warnings.warn(\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": "3252d85d",
   "metadata": {
    "papermill": {
     "duration": 2.015138,
     "end_time": "2023-11-10T10:09:19.768234",
     "exception": false,
     "start_time": "2023-11-10T10:09:17.753096",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "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": "0ad12a6f",
   "metadata": {
    "papermill": {
     "duration": 0.166712,
     "end_time": "2023-11-10T10:09:19.939479",
     "exception": false,
     "start_time": "2023-11-10T10:09:19.772767",
     "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": "a6bf1dab",
   "metadata": {
    "papermill": {
     "duration": 524.272471,
     "end_time": "2023-11-10T10:18:04.216146",
     "exception": false,
     "start_time": "2023-11-10T10:09:19.943675",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1790, 1574, 0.8793296089385475)"
      ]
     },
     "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": "12524d70",
   "metadata": {
    "papermill": {
     "duration": 0.004077,
     "end_time": "2023-11-10T10:18:04.225190",
     "exception": false,
     "start_time": "2023-11-10T10:18:04.221113",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cb003b6b",
   "metadata": {
    "papermill": {
     "duration": 0.01907,
     "end_time": "2023-11-10T10:18:04.248267",
     "exception": false,
     "start_time": "2023-11-10T10:18:04.229197",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('pkshatech/GLuCoSE-base-ja-v2',\n",
       " 0.808965162015302,\n",
       " 0.8490230332181042,\n",
       " 0.8793296089385475)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_id, jsts_score, jsick_score, miracl_recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2666bf46",
   "metadata": {
    "papermill": {
     "duration": 0.025368,
     "end_time": "2023-11-10T10:18:04.278514",
     "exception": false,
     "start_time": "2023-11-10T10:18:04.253146",
     "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.11.6"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 545.573896,
   "end_time": "2023-11-10T10:18:06.354125",
   "environment_variables": {},
   "exception": null,
   "input_path": "template_sentence-transformers.ipynb",
   "output_path": "17_intfloat_multilingual-e5-small.ipynb",
   "parameters": {
    "model_id": "intfloat/multilingual-e5-small",
    "passage_prefix": "passage: ",
    "query_prefix": "query: "
   },
   "start_time": "2023-11-10T10:09:00.780229",
   "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": {
     "026cdad54bb44ff5b30e70e3e2490ecf": {
      "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_fae6a3d2595f41ab877d989126a67584",
       "placeholder": "​",
       "style": "IPY_MODEL_860455510512472a881d9433f67fba70",
       "tabbable": null,
       "tooltip": null,
       "value": "100%"
      }
     },
     "2c35bbf2b01b4b8f871d75b71015d10f": {
      "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_2ceb0864f3774a069e413da95a7216a3",
       "max": 1,
       "min": 0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_e5a37620b26d42e59ea1f78ad86fa7ed",
       "tabbable": null,
       "tooltip": null,
       "value": 1
      }
     },
     "2ceb0864f3774a069e413da95a7216a3": {
      "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
      }
     },
     "2e24e018d1cb4281a04b3fed026047d0": {
      "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
      }
     },
     "3abe4e3607f749e692befe99d39225c2": {
      "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_026cdad54bb44ff5b30e70e3e2490ecf",
        "IPY_MODEL_2c35bbf2b01b4b8f871d75b71015d10f",
        "IPY_MODEL_ca3ab293e2804fa2babd6ded84e560d5"
       ],
       "layout": "IPY_MODEL_d2dfeda021d54a83a90a18d8919deab4",
       "tabbable": null,
       "tooltip": null
      }
     },
     "860455510512472a881d9433f67fba70": {
      "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
      }
     },
     "8ee01bc1f4bb458db8514a3051ce3e6d": {
      "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
      }
     },
     "ca3ab293e2804fa2babd6ded84e560d5": {
      "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_2e24e018d1cb4281a04b3fed026047d0",
       "placeholder": "​",
       "style": "IPY_MODEL_8ee01bc1f4bb458db8514a3051ce3e6d",
       "tabbable": null,
       "tooltip": null,
       "value": " 1/1 [00:00&lt;00:00,  2.60it/s]"
      }
     },
     "d2dfeda021d54a83a90a18d8919deab4": {
      "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
      }
     },
     "e5a37620b26d42e59ea1f78ad86fa7ed": {
      "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": ""
      }
     },
     "fae6a3d2595f41ab877d989126a67584": {
      "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
}