notebooks/209_bi_encoder_mse_loss.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "139bee68",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b893819e",
   "metadata": {},
   "source": [
    "# Train a bi-encoder to learn name-to-vec encodings\n",
    "Try phonemes, subwords, and/or n-grams using anchor-name pairs with MSELoss\n",
    "\n",
    "**Deprecated in favor of 208_bi_encoder_margin_mse_loss**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0364890",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "\n",
    "import joblib\n",
    "import pandas as pd\n",
    "from phonemizer.separator import Separator\n",
    "from phonemizer.backend import EspeakBackend\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm.auto import tqdm\n",
    "import ray\n",
    "from ray import air, tune\n",
    "from ray.air import session\n",
    "from ray.tune.search.hyperopt import HyperOptSearch\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "from src.data.filesystem import fopen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "774de195",
   "metadata": {},
   "outputs": [],
   "source": [
    "given_surname = \"given\"\n",
    "\n",
    "use_bigrams = True\n",
    "embedding_dim = 16\n",
    "max_tokens = 15\n",
    "learning_rate = 0.001\n",
    "batch_size = 32\n",
    "report_size = 10000\n",
    "num_epochs = 10\n",
    "\n",
    "triplets_path=f\"s3://familysearch-names/processed/tree-hr-{given_surname}-triplets.csv.gz\"\n",
    "tfidf_path=f\"s3://nama-data/data/models/fs-{given_surname}-tfidf-v2.joblib\"\n",
    "vocab_path = f\"s3://nama-data/data/models/fs-{given_surname}-espeak_phoneme_vocab.json\"\n",
    "bigrams_vocab_path = f\"s3://nama-data/data/models/fs-{given_surname}-espeak_phoneme_vocab_bigrams.json\"\n",
    "model_path = f\"../data/models/bi_encoder-{given_surname}\""
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "id": "c5ed943c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.shutdown()\n",
    "ray.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7c030b7",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a60fb5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read triplets\n",
    "triplets_df = pd.read_csv(triplets_path, na_filter=False)\n",
    "print(len(triplets_df))\n",
    "triplets_df.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72b99074",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read phoneme vocab\n",
    "with fopen(bigrams_vocab_path if use_bigrams else vocab_path, 'r') as f:\n",
    "    phoneme_vocab = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0429f1bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "phoneme_vocab['[UNK]'] = len(phoneme_vocab)\n",
    "phoneme_vocab['[PAD]'] = len(phoneme_vocab)\n",
    "vocab_size = len(phoneme_vocab)\n",
    "vocab_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2cc687d",
   "metadata": {},
   "source": [
    "## Set up generators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea7d66f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "espeak = EspeakBackend('en-us')\n",
    "separator = Separator(phone=' ', syllable=None, word='|')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6a6eec",
   "metadata": {},
   "source": [
    "## Create training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c39b91",
   "metadata": {},
   "outputs": [],
   "source": [
    "def tokenize(name, max_tokens):\n",
    "    tokens = [phoneme_vocab['[PAD]']] * max_tokens\n",
    "    unk = phoneme_vocab['[UNK]']\n",
    "    phonemes = espeak.phonemize([name], separator=separator, strip=True)[0].split(' ')\n",
    "    context_phoneme = 'START'\n",
    "    if use_bigrams:\n",
    "        phonemes.append('END')\n",
    "    for ix, phoneme in enumerate(phonemes):\n",
    "        if ix == max_tokens:\n",
    "            break\n",
    "        if use_bigrams:\n",
    "            phoneme_bigram = f\"{context_phoneme},{phoneme}\"\n",
    "            if phoneme_bigram in phoneme_vocab:\n",
    "                tokens[ix] = phoneme_vocab[phoneme_bigram]\n",
    "            else:\n",
    "                tokens[ix] = phoneme_vocab.get(phoneme, unk)\n",
    "        else:\n",
    "            tokens[ix] = phoneme_vocab.get(phoneme, unk)\n",
    "        context_phoneme = phoneme\n",
    "    return tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6655e260",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_anchors = triplets_df['anchor'].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f021d010",
   "metadata": {},
   "outputs": [],
   "source": [
    "# array of (anchor_tokens, name, target)\n",
    "all_data = []\n",
    "for tup in tqdm(triplets_df.itertuples()):\n",
    "    anchor_tokens = tokenize(tup.anchor, max_tokens)\n",
    "    pos_tokens = tokenize(tup.positive, max_tokens)\n",
    "    neg_tokens = tokenize(tup.negative, max_tokens)\n",
    "    easy_neg = random.choice(all_anchors)\n",
    "    easy_neg_tokens = tokenize(easy_neg, max_tokens)\n",
    "    # anchor, positive\n",
    "    all_data.append({\n",
    "        'anchor': torch.tensor(anchor_tokens),\n",
    "        'name': torch.tensor(pos_tokens),\n",
    "        'target': torch.tensor(tup.positive_score, dtype=torch.float),\n",
    "    })\n",
    "    # anchor, hard-negative\n",
    "    all_data.append({\n",
    "        'anchor': torch.tensor(anchor_tokens),\n",
    "        'name': torch.tensor(neg_tokens),\n",
    "        'target': torch.tensor(tup.negative_score, dtype=torch.float),\n",
    "    })\n",
    "    # anchor, easy-negative\n",
    "    if anchor_tokens == easy_neg_tokens:\n",
    "        continue\n",
    "    all_data.append({\n",
    "        'anchor': torch.tensor(anchor_tokens),\n",
    "        'name': torch.tensor(easy_neg_tokens),\n",
    "        'target': torch.tensor(0.0, dtype=torch.float)\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "547a4406",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data, val_data = train_test_split(all_data, test_size=0.10)\n",
    "print(len(train_data), len(val_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26b5e413",
   "metadata": {},
   "source": [
    "## Train bi-encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2603396e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_fn(anchors, names, labels):\n",
    "    anchor_name_pred_sim = (anchors * names).sum(dim=-1)\n",
    "    return F.mse_loss(anchor_name_pred_sim, labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97fbe744",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your bi-encoder model\n",
    "class BiEncoder(nn.Module):\n",
    "    def __init__(self, embedding_dim, vocab_size, max_tokens, pad_token):\n",
    "        super(BiEncoder, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.max_tokens = max_tokens\n",
    "        self.pad_token = pad_token\n",
    "        self.embedding = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)\n",
    "        self.forward_positional_embedding = nn.Embedding(num_embeddings=max_tokens+1, embedding_dim=embedding_dim)\n",
    "        self.backward_positional_embedding = nn.Embedding(num_embeddings=max_tokens+1, embedding_dim=embedding_dim)\n",
    "        self.pooling = nn.AdaptiveAvgPool1d(1)  # Pooling layer to create a single vector\n",
    "\n",
    "    def forward(self, input):\n",
    "        # get token embedding\n",
    "        embedded = self.embedding(input)  # Shape: (batch_size, max_tokens, embedding_dim)\n",
    "        # get forward positional embedding: pad token is position 0\n",
    "        positions = torch.arange(start=1, end=self.max_tokens+1).repeat(input.shape[0], 1)\n",
    "        forward_positions = torch.where(input == self.pad_token, 0, positions)\n",
    "        forward_positional_embedded = self.forward_positional_embedding(forward_positions)\n",
    "        # get backward positional embedding\n",
    "        backward_positions = torch.where(input == self.pad_token, 0, 1)\n",
    "        backward_n_tokens = backward_positions.sum(dim=1)\n",
    "        for ix in range(backward_n_tokens.shape[0]):\n",
    "            n_tokens = backward_n_tokens[ix]\n",
    "            backward = torch.arange(start=n_tokens, end=0, step=-1)\n",
    "            backward_positions[ix][:n_tokens] = backward\n",
    "        backward_positional_embedded = self.backward_positional_embedding(backward_positions)\n",
    "        # multiply embeddings\n",
    "        embedded = embedded * forward_positional_embedded * backward_positional_embedded\n",
    "        pooled = self.pooling(embedded.permute(0, 2, 1)).squeeze(2)  # Shape: (batch_size, embedding_dim)\n",
    "        return pooled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125b6c3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "def train(model, train_loader, val_loader, loss_fn, optimizer, num_epochs, verbose=True):\n",
    "    for epoch in range(num_epochs):\n",
    "        # make sure gradient tracking is on\n",
    "        model.train()\n",
    "        running_loss = 0\n",
    "\n",
    "        for ix, data in enumerate(train_loader):\n",
    "            # get batch\n",
    "            anchors = data['anchor']\n",
    "            names = data['name']\n",
    "            targets = data['target']\n",
    "\n",
    "            # zero gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # Forward pass\n",
    "            anchor_embeddings = model(anchors)  # Shape: (batch_size, embedding_dim)\n",
    "            name_embeddings = model(names)  # Shape: (batch_size, embedding_dim)\n",
    "\n",
    "            # Calculate loss\n",
    "            loss = loss_fn(anchor_embeddings, name_embeddings, targets)\n",
    "\n",
    "            # Backward pass and optimization step\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            # Calculate loss and report\n",
    "            if verbose:\n",
    "                running_loss += loss.item()\n",
    "                if ix % report_size == report_size - 1:\n",
    "                    avg_loss = running_loss / report_size  # loss per batch\n",
    "                    print(f\"Epoch {epoch} batch {ix} loss {avg_loss}\")\n",
    "                    running_loss = 0\n",
    "\n",
    "        # set model to evaluation mode\n",
    "        model.eval()\n",
    "\n",
    "        # disable gradient computation\n",
    "        running_loss = 0\n",
    "        num_val_batches = 0\n",
    "        with torch.no_grad():\n",
    "            for data in val_loader:\n",
    "                anchors = data['anchor']\n",
    "                names = data['name']\n",
    "                targets = data['target']\n",
    "                anchor_embeddings = model(anchors)  # Shape: (batch_size, embedding_dim)\n",
    "                name_embeddings = model(names)  # Shape: (batch_size, embedding_dim)\n",
    "                loss = loss_fn(anchor_embeddings, name_embeddings, targets)\n",
    "                running_loss += loss.item()  \n",
    "                num_val_batches += 1\n",
    "\n",
    "        # calculate average validation loss\n",
    "        val_loss = running_loss / num_val_batches\n",
    "        if verbose:\n",
    "            print(f\"VALIDATION: Epoch {epoch} loss {val_loss}\")\n",
    "        # epoch_model_path = f\"{model_path}-{epoch}\"\n",
    "        # torch.save(model.state_dict, epoch_model_path)\n",
    "        \n",
    "    # return final epoch validation loss\n",
    "    return val_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10fefd66",
   "metadata": {},
   "source": [
    "## Hyperparameter search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f7bcdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainer(config, train_data, val_data, pad_token):\n",
    "    learning_rate = config['learning_rate']\n",
    "    batch_size = config['batch_size']\n",
    "    embedding_dim = config['embedding_dim']\n",
    "    num_epochs = config['num_epochs']\n",
    "\n",
    "    # Create an instance of the bi-encoder model\n",
    "    model = BiEncoder(embedding_dim, vocab_size, max_tokens, pad_token)\n",
    "    # device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    # model.to(device)\n",
    "\n",
    "    # Define the optimizer\n",
    "    optimizer = optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "    # Create data loader\n",
    "    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "    val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=True) \n",
    "\n",
    "    val_loss = train(model, train_loader, val_loader, loss_fn, optimizer, num_epochs, verbose=False)\n",
    "\n",
    "    # Report the metrics to Ray\n",
    "    session.report({'loss': val_loss})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dd7db90",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_space={\n",
    "    \"learning_rate\": tune.loguniform(1e-4, 1e-2),\n",
    "    \"batch_size\": tune.choice([8,16,32,64]),\n",
    "    \"embedding_dim\": tune.choice([8,16,32,64]),\n",
    "    \"num_epochs\": tune.choice([5, 10, 20]),\n",
    "}\n",
    "\n",
    "starting_parameters = [{\n",
    "    \"learning_rate\": learning_rate,\n",
    "    \"batch_size\": batch_size,\n",
    "    \"embedding_dim\": embedding_dim,\n",
    "    \"num_epochs\": num_epochs,\n",
    "}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d35f339f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://docs.ray.io/en/latest/tune/api/doc/ray.tune.Tuner.html#ray.tune.Tuner\n",
    "\n",
    "search_alg = HyperOptSearch(points_to_evaluate=starting_parameters)\n",
    "\n",
    "callbacks = []\n",
    "# if wandb_api_key_file:\n",
    "#     callbacks.append(WandbLoggerCallback(\n",
    "#         project=\"nama\",\n",
    "#         entity=\"nama\",\n",
    "#         group=\"80_cluster_tune_\"+given_surname+\"_agglomerative\",\n",
    "#         notes=\"\",\n",
    "#         config=config._asdict(),\n",
    "#         api_key_file=wandb_api_key_file\n",
    "#     ))\n",
    "\n",
    "tuner = tune.Tuner(\n",
    "    tune.with_parameters(\n",
    "        trainer,\n",
    "        train_data=train_data,\n",
    "        val_data=val_data,\n",
    "        pad_token=phoneme_vocab['[PAD]']\n",
    "    ),\n",
    "    param_space=search_space,\n",
    "    tune_config=tune.TuneConfig(\n",
    "        mode='min',\n",
    "        metric='loss',\n",
    "        search_alg=search_alg,\n",
    "        num_samples=200,\n",
    "        max_concurrent_trials=2,\n",
    "        time_budget_s=8*60*60,\n",
    "    ),\n",
    "    run_config=air.RunConfig(\n",
    "        callbacks=callbacks,\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47e457b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = tuner.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cfd2d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# All results as pandas dataframe\n",
    "df = results.get_dataframe()\n",
    "print(len(df))\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2154f8a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get result that has the lowest validation loss\n",
    "best_result = results.get_best_result(metric='loss', mode='min', scope='all')\n",
    "\n",
    "# Parameters with the lowest loss\n",
    "best_result.metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dba77e31",
   "metadata": {},
   "source": [
    "## Review predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b50b883",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "embedding_dim = 32\n",
    "learning_rate = 0.0010\n",
    "num_epochs = 10\n",
    "# batch_size = best_result.config['batch_size']\n",
    "# learning_rate = best_result.config['learning_rate']\n",
    "# embedding_dim = best_result.config['embedding_dim']\n",
    "# num_epochs = best_result.config['num_epochs']\n",
    "\n",
    "# Create an instance of the bi-encoder model\n",
    "model = BiEncoder(embedding_dim, vocab_size, max_tokens, phoneme_vocab['[PAD]'])\n",
    "\n",
    "# Define the optimizer\n",
    "optimizer = optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "# Create data loader\n",
    "train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)\n",
    "val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "train(model, train_loader, val_loader, loss_fn, optimizer, num_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f70e637e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(name1, name2):\n",
    "    name1_tokens = tokenize(name1, max_tokens)\n",
    "    name2_tokens = tokenize(name2, max_tokens)\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        embeddings = model(torch.tensor([name1_tokens, name2_tokens]))\n",
    "    # return (embeddings[0] * embeddings[1]).sum(dim=-1).item()\n",
    "    return F.cosine_similarity(embeddings[0], embeddings[1], dim=-1).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e17eedd",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_correct = 0\n",
    "num_error = 0\n",
    "for row in triplets_df.sample(1000).itertuples():\n",
    "    anchor = row.anchor\n",
    "    pos = row.positive\n",
    "    neg = row.negative\n",
    "    pos_predict = predict(anchor, pos)\n",
    "    neg_predict = predict(anchor, neg)\n",
    "    correct = pos_predict > neg_predict\n",
    "    if correct:\n",
    "        num_correct += 1\n",
    "    else:\n",
    "        num_error += 1\n",
    "    if pos_predict < 0.6 or not correct:\n",
    "        print(anchor, pos, pos_predict, row.positive_score, '' if correct else 'ERROR')\n",
    "        print(anchor, neg, neg_predict, row.negative_score)\n",
    "print(f\"num_correct={num_correct} num_error={num_error}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf0c628c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for row in triplets_df.sample(100).itertuples():\n",
    "    anchor = row.anchor\n",
    "    pos = 'john'\n",
    "    sim = predict(anchor, pos)\n",
    "    if sim > 0.5:\n",
    "        print(anchor, pos, sim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "367d2e03",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47edace3",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}