notebooks/251_save_bi_encoder_weights.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "139bee68",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "028823c5",
   "metadata": {},
   "source": [
    "# Save bi-encoder model weights\n",
    "\n",
    "Save the weights for each token and position so we can use them in the java fs-nama."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0364890",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import torch\n",
    "\n",
    "from src.models.biencoder import BiEncoder\n",
    "from src.models.tokenizer import get_tokenize_function_and_vocab"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "774de195",
   "metadata": {},
   "outputs": [],
   "source": [
    "given_surname = \"given\"\n",
    "model_type = 'cecommon+0+aug-0-1'\n",
    "model_path = f\"../data/models/bi_encoder-{given_surname}-{model_type}.pth\"\n",
    "max_tokens = 10\n",
    "subwords_path=f\"../data/models/fs-{given_surname}-subword-tokenizer-2000f.json\"\n",
    "\n",
    "weights_path=f\"../data/models/bi_encoder-{given_surname}-{model_type}-weights.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34474d97",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.is_available())\n",
    "print(\"cuda total\", torch.cuda.get_device_properties(0).total_memory)\n",
    "print(\"cuda reserved\", torch.cuda.memory_reserved(0))\n",
    "print(\"cuda allocated\", torch.cuda.memory_allocated(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c030b7",
   "metadata": {},
   "source": [
    "## Load bi-encoder and vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f906a8e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.load(model_path)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b65957dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenize, tokenizer_vocab = get_tokenize_function_and_vocab(\n",
    "    max_tokens=max_tokens,\n",
    "    subwords_path=subwords_path,\n",
    ")\n",
    "len(tokenizer_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e97e8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenid2token = {token_id: token for token, token_id in tokenizer_vocab.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae53f9a5",
   "metadata": {},
   "source": [
    "## Save weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15651b39",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_weights = {tokenid2token[ix]: embedding for ix, embedding in enumerate(model.embedding.weight.tolist())}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01e0077e",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(embedding_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b9473c",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(embedding_weights['a'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa6b470",
   "metadata": {},
   "outputs": [],
   "source": [
    "positional_weights = model.positional_embedding.weight.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27644c58",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(positional_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a69751b",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = {\n",
    "    \"tokens\": embedding_weights,\n",
    "    \"positions\": positional_weights,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e001759",
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecaa3a70",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(weights_path, 'w') as f:\n",
    "    json.dump(weights, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb7c8ef9",
   "metadata": {},
   "source": [
    "## Test similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aff52b7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "name1 = \"richard\"\n",
    "name2 = \"rickert\"\n",
    "tokens1 = tokenize(name1)\n",
    "tokens2 = tokenize(name2)\n",
    "sim = model.predict(tokens1, tokens2)\n",
    "print(sim)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nama",
   "language": "python",
   "name": "nama"
  },
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}