notebooks/204_generate_subword_tokenizer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "235a10ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abcd12ed",
   "metadata": {},
   "source": [
    "# Generate subword tokenizer\n",
    "\n",
    "Use the training data from notebook 100 or preferred tree names to train a subword tokenizer.\n",
    "\n",
    "This notebook generates 4 different tokenizers. We ended up using just the base tokenizer - the one saved in tokenizer_path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4abac93",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import json\n",
    "import random\n",
    "from typing import List\n",
    "\n",
    "import boto3\n",
    "import Levenshtein\n",
    "import pandas as pd\n",
    "from tokenizers import models, Tokenizer, trainers, NormalizedString, PreTokenizedString\n",
    "from tokenizers.normalizers import BertNormalizer\n",
    "from tokenizers.pre_tokenizers import PreTokenizer, Whitespace\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from src.data.filesystem import fopen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4836cd31",
   "metadata": {},
   "outputs": [],
   "source": [
    "given_surname = 'surname'\n",
    "\n",
    "# run with 500, 1000, 1500, 1800, 2000\n",
    "vocab_size = 2000\n",
    "vocab_type = 'f'\n",
    "\n",
    "triplets_path=f\"../data/processed/tree-hr-{given_surname}-triplets-v2-1000.csv.gz\"\n",
    "pref_path = f\"s3://familysearch-names/processed/tree-preferred-{given_surname}-aggr.csv.gz\"\n",
    "train_path = f\"s3://familysearch-names/processed/tree-hr-{given_surname}-train-v2.csv.gz\"\n",
    "nama_bucket = 'nama-data'\n",
    "\n",
    "tokenizer_path=f\"../data/models/fs-{given_surname}-subword-tokenizer-{vocab_size}{vocab_type}.json\"\n",
    "edit_tokenizer_path=f\"../data/models/fs-{given_surname}-edit-subword-tokenizer-{vocab_size}.json\"\n",
    "tokenizer_bigrams_vocab_path = f\"../data/models/fs-{given_surname}-tokenizer_vocab_bigrams-{vocab_size}{vocab_type}.json\"\n",
    "edit_tokenizer_bigrams_vocab_path = f\"../data/models/fs-{given_surname}-edit_tokenizer_vocab_bigrams-{vocab_size}.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dc04ddf",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6b8c9e",
   "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": "9dcc16ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_names_df = pd.read_csv(train_path, na_filter=False)\n",
    "print(train_names_df.shape)\n",
    "train_names_df.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "007ace20",
   "metadata": {},
   "outputs": [],
   "source": [
    "pref_df = pd.read_csv(pref_path, na_filter=False)\n",
    "print(len(pref_df))\n",
    "print(pref_df['frequency'].sum())\n",
    "pref_df.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c596ddd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "if vocab_type == 'p':  # use preferred names\n",
    "    all_names = []\n",
    "    for row in tqdm(pref_df.to_dict('records')):\n",
    "        for _ in range(row['frequency']):\n",
    "            all_names.append(row['name'])\n",
    "    random.shuffle(all_names)\n",
    "elif vocab_type == 'f':  # use train names with frequency\n",
    "    all_names = []\n",
    "    for row in tqdm(train_names_df.to_dict('records')):\n",
    "        for _ in range(row['frequency']):\n",
    "            all_names.append(row['tree_name'])\n",
    "            all_names.append(row['record_name'])\n",
    "    random.shuffle(all_names)    \n",
    "else:  # use train names as a set\n",
    "    all_names = list(set(train_names_df['tree_name']) | set(train_names_df['record_name']))\n",
    "print(len(all_names))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7613698",
   "metadata": {},
   "source": [
    "## Calculate edit pieces based on anchor-positive pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f00784a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_edit_pieces(src, tar):\n",
    "    src_pieces = []\n",
    "    tar_pieces = []\n",
    "    opcodes = Levenshtein.opcodes(src, tar)\n",
    "    for (opcode, src_start, src_end, tar_start, tar_end) in opcodes:\n",
    "        if opcode == 'equal':\n",
    "            src_pieces.append(src[src_start:src_end])\n",
    "            tar_pieces.append(tar[tar_start:tar_end])\n",
    "        elif opcode == 'delete':\n",
    "            src_pieces.append(src[src_start:src_end])\n",
    "        elif opcode == 'insert':\n",
    "            tar_pieces.append(tar[tar_start:tar_end])\n",
    "        elif opcode == 'replace':\n",
    "            src_pieces.append(src[src_start:src_end])\n",
    "            tar_pieces.append(tar[tar_start:tar_end])\n",
    "        else:\n",
    "            print('Unexpected opcode', opcode)\n",
    "    return ','.join(src_pieces), ','.join(tar_pieces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0446886c",
   "metadata": {},
   "outputs": [],
   "source": [
    "anchor_pos_df = triplets_df[['anchor', 'positive']].drop_duplicates()\n",
    "len(anchor_pos_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a05ace17",
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_names = set()\n",
    "for src, tar in tqdm(zip(anchor_pos_df['anchor'], anchor_pos_df['positive'])):\n",
    "    src, tar = generate_edit_pieces(src, tar)\n",
    "    edit_names.add(src)\n",
    "    edit_names.add(tar)\n",
    "len(edit_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a828f394",
   "metadata": {},
   "source": [
    "## Pre-tokenize by splitting on edit pieces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f1ac392",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EditPiecePreTokenizer:\n",
    "    def split(self, i: int, normalized_string: NormalizedString) -> List[NormalizedString]:\n",
    "        # we need to call `str(normalized_string)` because split expects a str,\n",
    "        # not a NormalizedString\n",
    "        return [NormalizedString(s) for s in str(normalized_string).split(',')]\n",
    "    \n",
    "    def pre_tokenize(self, pretok: PreTokenizedString):\n",
    "        return pretok.split(self.split)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e8ef21e",
   "metadata": {},
   "source": [
    "## Generate Subwords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe6cbd3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "special_tokens = [\"[UNK]\", \"[PAD]\", \"[CLS]\", \"[SEP]\", \"[MASK]\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94362d0a",
   "metadata": {},
   "source": [
    "### from edit pieces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbfdac30",
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_tokenizer = Tokenizer(models.WordPiece(unk_token=\"[UNK]\"))\n",
    "edit_tokenizer.pre_tokenizer = PreTokenizer.custom(EditPiecePreTokenizer())\n",
    "\n",
    "trainer = trainers.WordPieceTrainer(vocab_size=vocab_size, special_tokens=special_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99bce600",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train tokenizer from edit pieces\n",
    "def get_edit_names():\n",
    "    for name in edit_names:\n",
    "        yield name\n",
    "        \n",
    "edit_tokenizer.train_from_iterator(get_edit_names(), trainer=trainer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d13b63",
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_tokenizer.get_vocab()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cbcaab3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# now that the tokenizer has been trained, we don't need the pre-tokenizer any more\n",
    "# so default it to whitespace\n",
    "edit_tokenizer.pre_tokenizer = Whitespace()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d420327a",
   "metadata": {},
   "source": [
    "## Generate subwords from all names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "365b8fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(models.WordPiece(unk_token=\"[UNK]\"))\n",
    "\n",
    "trainer = trainers.WordPieceTrainer(vocab_size=vocab_size, special_tokens=special_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c15bdf95",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(all_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32bb97b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# train tokenizer from all names\n",
    "def get_all_names_batch(batch_size=10000):\n",
    "    for i in range(0, len(all_names), batch_size):\n",
    "        yield all_names[i : i + batch_size]\n",
    "        \n",
    "tokenizer.train_from_iterator(get_all_names_batch(), \n",
    "                              trainer=trainer, \n",
    "                              length=len(all_names))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3512862f",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.get_vocab()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fe3fe08",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(tokenizer.get_vocab())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "715bb28a",
   "metadata": {},
   "source": [
    "## Review sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47412abe",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 100\n",
    "sample_df = triplets_df.sample(sample_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dee5928f",
   "metadata": {},
   "outputs": [],
   "source": [
    "for anchor, positive in zip(sample_df['anchor'], sample_df['positive']):\n",
    "    print(anchor, positive)\n",
    "    print('edit', edit_tokenizer.encode(anchor).tokens, edit_tokenizer.encode(positive).tokens)\n",
    "    print(' all', tokenizer.encode(anchor).tokens, tokenizer.encode(positive).tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27b74146",
   "metadata": {},
   "source": [
    "## Save tokenizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25c36633",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.save(tokenizer_path)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51a2754c",
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_tokenizer.save(edit_tokenizer_path)  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8dd90e5",
   "metadata": {},
   "source": [
    "## Calculate subwords, subword-bigrams, and lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4b418ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "subword_counter = Counter()\n",
    "subword_bigrams_counter = Counter()\n",
    "subword_lengths = Counter()\n",
    "edit_subword_counter = Counter()\n",
    "edit_subword_bigrams_counter = Counter()\n",
    "edit_subword_lengths = Counter()\n",
    "tokenizer_cache = {}\n",
    "edit_tokenizer_cache = {}\n",
    "\n",
    "sample_size = 1000_000\n",
    "\n",
    "for name in tqdm(all_names[:sample_size]):\n",
    "    if name not in tokenizer_cache:\n",
    "        subwords = tokenizer.encode(name).tokens\n",
    "        tokenizer_cache[name] = subwords\n",
    "    else:\n",
    "        subwords = tokenizer_cache[name]\n",
    "        \n",
    "    for subword in subwords:\n",
    "        subword_counter[subword] += 1\n",
    "    context_subword = 'START'\n",
    "    subwords.append('END')\n",
    "    for subword in subwords:\n",
    "        subword_bigrams_counter[f\"{context_subword},{subword}\"] += 1\n",
    "        context_subword = subword\n",
    "    subword_lengths[len(subwords)] += 1\n",
    "    \n",
    "    if name not in edit_tokenizer_cache:\n",
    "        subwords = edit_tokenizer.encode(name).tokens\n",
    "        edit_tokenizer_cache[name] = subwords\n",
    "    else:\n",
    "        subwords = edit_tokenizer_cache[name]\n",
    "        \n",
    "    for subword in subwords:\n",
    "        edit_subword_counter[subword] += 1\n",
    "    context_subword = 'START'\n",
    "    subwords.append('END')\n",
    "    for subword in subwords:\n",
    "        edit_subword_bigrams_counter[f\"{context_subword},{subword}\"] += 1\n",
    "        context_subword = subword\n",
    "    edit_subword_lengths[len(subwords)] += 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d08687eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "subword_counter.most_common()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d17de5fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_subword_counter.most_common()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1946445b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(subword_bigrams_counter))\n",
    "subword_bigrams_counter.most_common(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b694ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(edit_subword_bigrams_counter))\n",
    "edit_subword_bigrams_counter.most_common(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0608c175",
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_subword_lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df2ba267",
   "metadata": {},
   "outputs": [],
   "source": [
    "subword_lengths"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0886802",
   "metadata": {},
   "source": [
    "## Save bigram tokenizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10976d11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save tokenizer bigrams vocabulary\n",
    "tokenizer_bigrams_vocab = {}\n",
    "ix = 0\n",
    "for subword in tokenizer.get_vocab():\n",
    "    tokenizer_bigrams_vocab[subword] = ix\n",
    "    ix += 1\n",
    "for bigram, _ in subword_bigrams_counter.most_common(vocab_size):\n",
    "    tokenizer_bigrams_vocab[bigram] = ix\n",
    "    ix += 1\n",
    "\n",
    "print(len(tokenizer_bigrams_vocab))\n",
    "\n",
    "with fopen(tokenizer_bigrams_vocab_path, 'w') as f:\n",
    "    json.dump(tokenizer_bigrams_vocab, f)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b92b98a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save edit tokenizer bigrams vocabulary\n",
    "edit_tokenizer_bigrams_vocab = {}\n",
    "ix = 0\n",
    "for subword in edit_tokenizer.get_vocab():\n",
    "    edit_tokenizer_bigrams_vocab[subword] = ix\n",
    "    ix += 1\n",
    "for bigram, _ in edit_subword_bigrams_counter.most_common(vocab_size):\n",
    "    edit_tokenizer_bigrams_vocab[bigram] = ix\n",
    "    ix += 1\n",
    "    \n",
    "print(len(edit_tokenizer_bigrams_vocab))\n",
    "    \n",
    "with fopen(edit_tokenizer_bigrams_vocab_path, 'w') as f:\n",
    "    json.dump(edit_tokenizer_bigrams_vocab, f)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd58bf97",
   "metadata": {},
   "source": [
    "## Test load tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47f8370a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03222467",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a497f58c",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in ['chesworth', 'quass']:\n",
    "    print(loaded_tokenizer.convert_ids_to_tokens(loaded_tokenizer.encode(name)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad4cd4a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "alt_tokenizer_path = '../data/models/fs-given-subword-tokenizer-2000f.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21ec9f5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "alt_loaded_tokenizer = PreTrainedTokenizerFast(tokenizer_file=alt_tokenizer_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f0c3bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "alt_loaded_tokenizer.convert_ids_to_tokens(alt_loaded_tokenizer.encode('zacharias'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c630e49c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_size = 20\n",
    "sample_df = triplets_df.sample(sample_size)\n",
    "\n",
    "for anchor in sample_df['anchor']:\n",
    "    tokens = loaded_tokenizer.convert_ids_to_tokens(loaded_tokenizer.encode(anchor))\n",
    "    alt_tokens = alt_loaded_tokenizer.convert_ids_to_tokens(alt_loaded_tokenizer.encode(anchor))\n",
    "    if tokens != alt_tokens:\n",
    "        print(anchor, tokens, alt_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59a9d7c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_tokenizer.convert_ids_to_tokens(loaded_tokenizer.encode('jewel'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4140e96d",
   "metadata": {},
   "outputs": [],
   "source": [
    "alt_loaded_tokenizer.convert_ids_to_tokens(alt_loaded_tokenizer.encode('jewel'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c7ffcdb",
   "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
}