30_hotchpotch_japanse_splade.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.949996Z",
     "start_time": "2023-10-07T08:50:44.497041Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import os\n",
    "\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hotchpotch/miniconda3/envs/cross-encoder-japanese/lib/python3.11/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:11: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from tqdm.autonotebook import tqdm, trange\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import CrossEncoder\n",
    "import torch\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-10-07T08:50:45.965421Z",
     "start_time": "2023-10-07T08:50:45.951998Z"
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "miracle_n_hard_negs = 300\n",
    "miracle_n_recall = 30"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters\n",
    "model_id = \"hotchpotch/japanese-splade-base-v1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForMaskedLM(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(32768, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0-11): 12 x BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSdpaSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "            (intermediate_act_fn): GELUActivation()\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (cls): BertOnlyMLMHead(\n",
       "    (predictions): BertLMPredictionHead(\n",
       "      (transform): BertPredictionHeadTransform(\n",
       "        (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "        (transform_act_fn): GELUActivation()\n",
       "        (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      )\n",
       "      (decoder): Linear(in_features=768, out_features=32768, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForMaskedLM, AutoTokenizer\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "model_name = model_id\n",
    "\n",
    "model = AutoModelForMaskedLM.from_pretrained(model_name)\n",
    "if device == \"cuda\":\n",
    "    model.to(device)\n",
    "    model.half()  # faster inference\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def splade_max_pooling(logits, attention_mask):\n",
    "    relu_log = torch.log(1 + torch.relu(logits))\n",
    "    weighted_log = relu_log * attention_mask.unsqueeze(-1)\n",
    "    max_val, _ = torch.max(weighted_log, dim=1)\n",
    "    return max_val\n",
    "\n",
    "\n",
    "def splade_vector(texts: list[str], bs=64):\n",
    "    pooled_all = []\n",
    "    for i in range(0, len(texts), bs):\n",
    "        with torch.no_grad():\n",
    "            batch_texts = texts[i : i + bs]\n",
    "            inputs = tokenizer(\n",
    "                batch_texts,\n",
    "                return_tensors=\"pt\",\n",
    "                padding=True,\n",
    "                truncation=True,\n",
    "                max_length=512,\n",
    "            )\n",
    "            inputs = {k: v.to(device) for k, v in inputs.items()}\n",
    "            attention_mask = inputs[\"attention_mask\"]\n",
    "            input_ids = inputs[\"input_ids\"]\n",
    "            outputs = model(input_ids, attention_mask=attention_mask)\n",
    "            logits = outputs.logits\n",
    "            pooled = splade_max_pooling(logits, attention_mask)\n",
    "            pooled_all.append(pooled)\n",
    "\n",
    "    return torch.cat(pooled_all).cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(32768,)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 1000個のテキストをランダムに生成\n",
    "import random\n",
    "\n",
    "random.seed(42)\n",
    "texts = [\n",
    "    \"\".join([random.choice(\"abcdefghijklmnopqrstuvwxyz\") for _ in range(100)])\n",
    "    for _ in range(1000)\n",
    "]\n",
    "\n",
    "# 1000個のテキストに対して、それぞれのテキストの特徴ベクトルを計算\n",
    "vectors = []\n",
    "for vec in splade_vector(texts):\n",
    "    vectors.append(vec)\n",
    "vectors[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Miracle\n",
    "\n",
    "- Need access token for huggingface\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "No module named 'dotenv'\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "try:\n",
    "    import dotenv\n",
    "\n",
    "    dotenv.load_dotenv(\"huggingface_access_token\", override=True)\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['query_id', 'query', 'positive_passages', 'negative_passages'],\n",
       "    num_rows: 860\n",
       "})"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import datasets\n",
    "\n",
    "# query and positives\n",
    "ds = datasets.load_dataset(\"miracl/miracl\", \"ja\", split=\"dev\")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['docid', 'title', 'text'],\n",
       "        num_rows: 6953614\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 10,
     "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": 11,
   "metadata": {},
   "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": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# hard negatives\n",
    "with open(\"./miracl_hard_negs_1000.json\") as f:\n",
    "    hn = json.loads(f.read())\n",
    "(\n",
    "    len(hn),\n",
    "    list(hn.keys())[:5],\n",
    "    hn[\"0\"].keys(),\n",
    "    hn[\"0\"][\"docids\"][:2],\n",
    "    hn[\"0\"][\"indices\"][:2],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "target_curpus_ids: 227726\n"
     ]
    }
   ],
   "source": [
    "target_curpus_ids = set()\n",
    "for item in ds:\n",
    "    query = item[\"query\"]\n",
    "    # passages are set(300 hard negatives + positives)\n",
    "    positive_docids = [pp[\"docid\"] for pp in item[\"positive_passages\"]]\n",
    "    hn_docids = hn[item[\"query_id\"]][\"docids\"][:miracle_n_hard_negs]\n",
    "    target_curpus_ids.update(positive_docids)\n",
    "    target_curpus_ids.update(hn_docids)\n",
    "print(f\"target_curpus_ids: {len(target_curpus_ids)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6953614/6953614 [01:14<00:00, 92802.37it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "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 tqdm(corpus[\"train\"])}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6953614"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(corpus_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "\n",
    "# for \"Be aware, overflowing tokens are not returned for the setting you have chosen\" warning\n",
    "transformers.logging.set_verbosity_error()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/860 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 860/860 [05:35<00:00,  2.56it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1790, 1656, 0.9251396648044693)"
      ]
     },
     "execution_count": 16,
     "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",
    "n_total_pos = 0\n",
    "n_total_tp = 0\n",
    "\n",
    "for item in tqdm(ds):\n",
    "    # query\n",
    "    # query_emb = model.encode([retrieve_query_prefix + item[\"query\"]])\n",
    "    query_emb = splade_vector([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(\n",
    "    #     [retrieve_passage_prefix + text for text in target_texts]\n",
    "    # )\n",
    "    target_embs = splade_vector(target_texts)\n",
    "\n",
    "    # topK\n",
    "    # topk_indices = np.argsort(cdist(query_emb, target_embs, metric=\"cosine\"))[0][\n",
    "    #     :miracle_n_recall\n",
    "    # ]\n",
    "    scores = np.dot(query_emb, target_embs.T)\n",
    "    topk_indices = np.argsort(-scores).tolist()[0][:miracle_n_recall]\n",
    "    n_tp = len(\n",
    "        set(topk_indices) & set(range(len(positive_docids)))\n",
    "    )  # positives are first indices\n",
    "    n_pos = len(positive_docids)\n",
    "\n",
    "    n_total_pos += n_pos\n",
    "    n_total_tp += n_tp\n",
    "\n",
    "\n",
    "miracl_recall = n_total_tp / n_total_pos\n",
    "\n",
    "n_total_pos, n_total_tp, miracl_recall"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('hotchpotch/japanese-splade-base-v1', None, None, 0.9251396648044693)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jsts_score = None\n",
    "jsick_score = None\n",
    "model_id, jsts_score, jsick_score, miracl_recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "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": "cross-encoder-japanese",
   "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.9"
  },
  "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
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}