notebooks/64_analyze_scores.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "896a97b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c9d8205",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Compare swivel and levenshtein scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad2e0f1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "\n",
    "import jellyfish\n",
    "import matplotlib.pyplot as plt\n",
    "from mpire import WorkerPool\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "from sklearn.model_selection import train_test_split\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "import wandb\n",
    "\n",
    "from src.data.filesystem import fopen\n",
    "from src.data.utils import load_dataset\n",
    "from src.eval.utils import similars_to_ndarray\n",
    "from src.models.swivel import SwivelModel, get_best_swivel_matches\n",
    "from src.models.utils import remove_padding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dfe0578",
   "metadata": {},
   "outputs": [],
   "source": [
    "# config\n",
    "\n",
    "given_surname = \"given\"\n",
    "vocab_size = 610000 if given_surname == \"given\" else 2100000\n",
    "sample_size = 5000\n",
    "embed_dim = 100\n",
    "encoder_layers = 2\n",
    "num_matches = 5000\n",
    "batch_size = 256\n",
    "Config = namedtuple(\"Config\", [\n",
    "    \"eval_path\",\n",
    "    \"embed_dim\",\n",
    "    \"swivel_vocab_path\",\n",
    "    \"swivel_model_path\",\n",
    "])\n",
    "config = Config(\n",
    "    eval_path=f\"s3://familysearch-names/processed/tree-hr-{given_surname}-train.csv.gz\",\n",
    "    embed_dim=embed_dim,\n",
    "    swivel_vocab_path=f\"s3://nama-data/data/models/fs-{given_surname}-swivel-vocab-{vocab_size}-augmented.csv\",\n",
    "    swivel_model_path=f\"s3://nama-data/data/models/fs-{given_surname}-swivel-model-{vocab_size}-{embed_dim}-augmented.pth\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c13e7bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# wandb.init(\n",
    "#     project=\"nama\",\n",
    "#     entity=\"nama\",\n",
    "#     name=\"64_analyze_scores\",\n",
    "#     group=given_surname,\n",
    "#     notes=\"\",\n",
    "#     config=config._asdict(),\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e17e3d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\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": "code",
   "execution_count": null,
   "id": "390f4a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_names_eval, weighted_actual_names_eval, candidate_names_eval = load_dataset(config.eval_path, is_eval=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e269e938",
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_df = pd.read_csv(fopen(config.swivel_vocab_path, \"rb\"))\n",
    "swivel_vocab = {name: _id for name, _id in zip(vocab_df[\"name\"], vocab_df[\"index\"])}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b3e296",
   "metadata": {},
   "outputs": [],
   "source": [
    "swivel_model = SwivelModel(len(swivel_vocab), config.embed_dim)\n",
    "swivel_model.load_state_dict(torch.load(fopen(config.swivel_model_path, \"rb\"), map_location=torch.device(device)))\n",
    "swivel_model.to(device)\n",
    "swivel_model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "737919ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "_, input_names_eval_sample, _, weighted_actual_names_eval_sample = \\\n",
    "    train_test_split(input_names_eval, weighted_actual_names_eval, test_size=sample_size)\n",
    "candidate_names_eval_sample = candidate_names_eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47c0b96e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"input_names_eval_sample\", len(input_names_eval_sample))\n",
    "print(\"weighted_actual_names_eval_sample\", len(weighted_actual_names_eval_sample))\n",
    "print(\"candidate_names_eval_sample\", len(candidate_names_eval_sample))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "721651a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_similarity_to(name):\n",
    "    name = remove_padding(name)\n",
    "\n",
    "    def calc_similarity(row):\n",
    "        cand_name = remove_padding(row[0])\n",
    "        dist = jellyfish.levenshtein_distance(name, cand_name)\n",
    "        return 1 - (dist / max(len(name), len(cand_name)))\n",
    "\n",
    "    return calc_similarity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cc40b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_similars(shared, names, _=None):\n",
    "    candidate_names_test, k = shared\n",
    "    \n",
    "    def get_similars_for_name(name):\n",
    "        scores = np.apply_along_axis(calc_similarity_to(name), 1, candidate_names_test[:, None])\n",
    "\n",
    "        # sorted_scores_idx = np.argsort(scores)[::-1][:k]\n",
    "        partitioned_idx = np.argpartition(scores, -k)[-k:]\n",
    "        sorted_partitioned_idx = np.argsort(scores[partitioned_idx])[::-1]\n",
    "        sorted_scores_idx = partitioned_idx[sorted_partitioned_idx]\n",
    "\n",
    "        candidate_names = candidate_names_test[sorted_scores_idx]\n",
    "        candidate_scores = scores[sorted_scores_idx]\n",
    "\n",
    "        return list(zip(candidate_names, candidate_scores))\n",
    "    \n",
    "    result = []\n",
    "    for name in names:\n",
    "        result.append(get_similars_for_name(name))\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02d09438",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_batches(names, batch_size):\n",
    "    batches = []\n",
    "    for ix in range(0, len(names), batch_size):\n",
    "        # batches are tuples to keep mpire from expanding the batch \n",
    "        batches.append((names[ix:ix + batch_size], ix))\n",
    "    return batches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02386fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "swivel_names_scores = get_best_swivel_matches(model=swivel_model, \n",
    "                                              vocab=swivel_vocab, \n",
    "                                              input_names=input_names_eval_sample,\n",
    "                                              candidate_names=candidate_names_eval_sample, \n",
    "                                              encoder_model=None,\n",
    "                                              k=num_matches, \n",
    "                                              batch_size=batch_size,\n",
    "                                              add_context=True,\n",
    "                                              n_jobs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a28f2fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sum(len(names_scores) for names_scores in swivel_names_scores))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ba1a35",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_names_batches = create_batches(input_names_eval_sample, batch_size=batch_size)\n",
    "with WorkerPool(\n",
    "    shared_objects=(candidate_names_eval_sample, num_matches),\n",
    ") as pool:\n",
    "    lev_names_scores = pool.map(get_similars, input_names_batches, progress_bar=True)\n",
    "# flatten\n",
    "lev_names_scores = [name_score for batch in lev_names_scores for name_score in batch]\n",
    "# convert to ndarray\n",
    "lev_names_scores = similars_to_ndarray(lev_names_scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8776c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sum(len(names_scores) for names_scores in lev_names_scores))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d0666f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# find pairs in both with score above a threshold\n",
    "swivel_threshold = 0.45\n",
    "lev_threshold = 0.55\n",
    "sample_rate = 0.01\n",
    "xs = []\n",
    "ys = []\n",
    "cs = []\n",
    "xs_pos = []\n",
    "ys_pos = []\n",
    "xs_neg = []\n",
    "ys_neg = []\n",
    "weights = []\n",
    "actual_score_counts = []\n",
    "swivel_score_counts = []\n",
    "lev_score_counts = []\n",
    "all_candidate_names = set(candidate_names_eval_sample)\n",
    "for input_name, wans, swivels, levs in \\\n",
    "    zip(input_names_eval_sample, weighted_actual_names_eval_sample, swivel_names_scores, lev_names_scores):\n",
    "    # actuals - ensure names are in all_candidate_names\n",
    "    actual_weights = {name: weight for name, weight, _ in wans if name in all_candidate_names}\n",
    "    actual_score_counts.append(len(actual_weights))\n",
    "    # swivel\n",
    "    swivel_scores = {name: score for name, score in swivels if score >= swivel_threshold}\n",
    "    swivel_names = set(swivel_scores.keys())\n",
    "    swivel_score_counts.append(len(swivel_scores))\n",
    "    # levenshtein\n",
    "    lev_scores = {name: score for name, score in levs if score >= lev_threshold}\n",
    "    lev_names = set(lev_scores.keys())\n",
    "    lev_score_counts.append(len(lev_scores))\n",
    "    \n",
    "    # count various scores\n",
    "    candidate_names = swivel_names.intersection(lev_names)\n",
    "    for candidate_name in candidate_names:\n",
    "        if random.random() > sample_rate:\n",
    "            continue\n",
    "        swivel_score = swivel_scores[candidate_name]\n",
    "        lev_score = lev_scores[candidate_name]\n",
    "        xs.append(swivel_score)\n",
    "        ys.append(lev_score)\n",
    "        if candidate_name in actual_weights:\n",
    "            cs.append('green')\n",
    "            xs_pos.append(swivel_score)\n",
    "            ys_pos.append(lev_score)\n",
    "            weights.append(actual_weights[candidate_name])\n",
    "            del actual_weights[candidate_name]\n",
    "        else:\n",
    "            cs.append('red')\n",
    "            xs_neg.append(swivel_score)\n",
    "            ys_neg.append(lev_score)\n",
    "#     for name in actual_weights.keys():\n",
    "#         if name not in swivel_names:\n",
    "#             print(\"swivel\", input_name, name)\n",
    "#         if name not in lev_names:\n",
    "#             print(\"lev\", input_name, name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7306b4b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sum(swivel_score_counts), sum(lev_score_counts))\n",
    "print(len(cs), len([c for c in cs if c == 'green']), sum(actual_score_counts)*sample_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "693b2d09",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure, ax = plt.subplots(1, 1, figsize=(20, 15))\n",
    "ax.set_title(\"Swivel vs Levenshtein score\")\n",
    "ax.scatter(x=xs, y=ys, c=cs)\n",
    "plt.xlabel(\"swivel score\")\n",
    "plt.ylabel(\"levenshtein score\")\n",
    "plt.xlim([swivel_threshold, 1.0])\n",
    "plt.ylim([lev_threshold, 1.0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0170f455",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure, ax = plt.subplots(1, 1, figsize=(20, 15))\n",
    "ax.set_title(\"Swivel vs Levenshtein score - positive only\")\n",
    "ax.scatter(x=xs_pos, y=ys_pos)\n",
    "plt.xlabel(\"swivel score\")\n",
    "plt.ylabel(\"levenshtein score\")\n",
    "plt.xlim([swivel_threshold, 1.0])\n",
    "plt.ylim([lev_threshold, 1.0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "335f7609",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure, ax = plt.subplots(1, 1, figsize=(20, 15))\n",
    "ax.set_title(\"Swivel vs Levenshtein score - negative only\")\n",
    "ax.scatter(x=xs_neg, y=ys_neg)\n",
    "plt.xlabel(\"swivel score\")\n",
    "plt.ylabel(\"levenshtein score\")\n",
    "plt.xlim([swivel_threshold, 1.0])\n",
    "plt.ylim([lev_threshold, 1.0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4a541af",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure, ax = plt.subplots(1, 1, figsize=(20, 15))\n",
    "ax.set_title(\"Weights\")\n",
    "ax.hist(x=weights, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e63b68",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure, ax = plt.subplots(1, 1, figsize=(20, 15))\n",
    "ax.set_title(\"Swivel score counts\")\n",
    "ax.hist(x=swivel_score_counts, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62b1ccc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure, ax = plt.subplots(1, 1, figsize=(20, 15))\n",
    "ax.set_title(\"Levenshtein score counts\")\n",
    "ax.hist(x=lev_score_counts, bins=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5259e6e0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}