examples/06_predicting_like_terms.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predicting Like Polynomial Terms\n",
    "\n",
    "Remember in Algebra how you had to combine \"like terms\" to simplify problems? \n",
    "\n",
    "You'd see expressions such as `60 + 2x^3 - 6x + x^3 + 17x` in which there are **5** total terms but only **4** are \"like terms\". \n",
    "\n",
    "`2x^3` and `x^3` are like, and `-6x` and `17x` are like, while `60` doesn't have any like siblings.\n",
    "\n",
    "Can we teach a model to predict that there are `4` like terms in the above expression?\n",
    "\n",
    "Let's give it a shot using [Mathy](https://mathy.ai) to generate math problems and [thinc](https://github.com/explosion/thinc) to build a regression model that outputs the number of like terms in each input problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"thinc>=8.0.0\" mathy_core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sketch a Model\n",
    "\n",
    "Before we get started it can be good to have an idea of what input/output shapes we want for our model.\n",
    "\n",
    "We'll convert text math problems into lists of lists of integers, so our example (X) type can be represented using thinc's `Ints2d` type.\n",
    "\n",
    "The model will predict how many like terms there are in each sequence, so our output (Y) type can represented with the `Floats2d` type.\n",
    "\n",
    "Knowing the thinc types we want enables us to create an alias for our model, so we only have to type out the verbose generic signature once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from thinc.api import Model\n",
    "from thinc.types import Ints2d, Floats1d\n",
    "\n",
    "ModelX = Ints2d\n",
    "ModelY = Floats1d\n",
    "ModelT = Model[List[ModelX], ModelY]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encode Text Inputs\n",
    "\n",
    "Mathy generates ascii-math problems and we have to encode them into integers that the model can process. \n",
    "\n",
    "To do this we'll build a vocabulary of all the possible characters we'll see, and map each input character to its index in the list.\n",
    "\n",
    "For math problems our vocabulary will include all the characters of the alphabet, numbers 0-9, and special characters like `*`, `-`, `.`, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from thinc.api import Model\n",
    "from thinc.types import Ints2d, Floats1d\n",
    "from thinc.api import Ops, get_current_ops\n",
    "\n",
    "vocab = \" .+-/^*()[]-01234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\"\n",
    "\n",
    "def encode_input(text: str) -> ModelX:\n",
    "    ops: Ops = get_current_ops()\n",
    "    indices: List[List[int]] = []\n",
    "    for c in text:\n",
    "        if c not in vocab:\n",
    "            raise ValueError(f\"'{c}' missing from vocabulary in text: {text}\")\n",
    "        indices.append([vocab.index(c)])\n",
    "    return ops.asarray2i(indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It\n",
    "\n",
    "Let's try it out on some fixed data to be sure it works. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = encode_input(\"4+2\")\n",
    "assert outputs[0][0] == vocab.index(\"4\")\n",
    "assert outputs[1][0] == vocab.index(\"+\")\n",
    "assert outputs[2][0] == vocab.index(\"2\")\n",
    "print(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Math Problems\n",
    "\n",
    "We'll use Mathy to generate random polynomial problems with a variable number of like terms. The generated problems will act as training data for our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Optional, Set\n",
    "import random\n",
    "from mathy_core.problems import gen_simplify_multiple_terms\n",
    "\n",
    "def generate_problems(number: int, exclude: Optional[Set[str]] = None) -> List[str]:\n",
    "    if exclude is None:\n",
    "        exclude = set()\n",
    "    problems: List[str] = []\n",
    "    while len(problems) < number:\n",
    "        text, complexity = gen_simplify_multiple_terms(\n",
    "            random.randint(2, 6),\n",
    "            noise_probability=1.0,\n",
    "            noise_terms=random.randint(2, 10),\n",
    "            op=[\"+\", \"-\"],\n",
    "        )\n",
    "        assert text not in exclude, \"duplicate problem generated!\"\n",
    "        exclude.add(text)\n",
    "        problems.append(text)\n",
    "    return problems"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_problems(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Count Like Terms\n",
    "\n",
    "Now that we can generate input problems, we'll need a function that can count the like terms in each one and return the value for use as a label.\n",
    "\n",
    "To accomplish this we'll use a few helpers from mathy to enumerate the terms and compare them to see if they're like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, List, Dict\n",
    "from mathy_core import MathExpression, ExpressionParser, get_terms, get_term_ex, TermEx\n",
    "from mathy_core.problems import mathy_term_string\n",
    "\n",
    "parser = ExpressionParser()\n",
    "\n",
    "def count_like_terms(input_problem: str) -> int:\n",
    "    expression: MathExpression = parser.parse(input_problem)\n",
    "    term_nodes: List[MathExpression] = get_terms(expression)\n",
    "    node_groups: Dict[str, List[MathExpression]] = {}\n",
    "    for term_node in term_nodes:\n",
    "        ex: Optional[TermEx] = get_term_ex(term_node)\n",
    "        assert ex is not None, f\"invalid expression {term_node}\"\n",
    "        key = mathy_term_string(variable=ex.variable, exponent=ex.exponent)\n",
    "        if key == \"\":\n",
    "            key = \"const\"\n",
    "        if key not in node_groups:\n",
    "            node_groups[key] = [term_node]\n",
    "        else:\n",
    "            node_groups[key].append(term_node)\n",
    "    like_terms = 0\n",
    "    for k, v in node_groups.items():\n",
    "        if len(v) <= 1:\n",
    "            continue\n",
    "        like_terms += len(v)\n",
    "    return like_terms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert count_like_terms(\"4x - 2y + q\") == 0\n",
    "assert count_like_terms(\"x + x + z\") == 2\n",
    "assert count_like_terms(\"4x + 2x - x + 7\") == 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Problem/Answer pairs\n",
    "\n",
    "Now that we can generate problems, count the number of like terms in them, and encode their text into integers, we have the pieces required to generate random problems and answers that we can train a neural network with.\n",
    "\n",
    "Let's write a function that will return a tuple of: the problem text, its encoded example form, and the output label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "from thinc.api import Ops, get_current_ops\n",
    "\n",
    "def to_example(input_problem: str) -> Tuple[str, ModelX, ModelY]:\n",
    "    ops: Ops = get_current_ops()\n",
    "    encoded_input = encode_input(input_problem)\n",
    "    like_terms = count_like_terms(input_problem)\n",
    "    return input_problem, encoded_input, ops.asarray1f([like_terms])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text, X, Y = to_example(\"x+2x\")\n",
    "assert text == \"x+2x\"\n",
    "assert X[0] == vocab.index(\"x\")\n",
    "assert Y[0] == 2\n",
    "print(text, X, Y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a Model\n",
    "\n",
    "Now that we can generate X/Y values, let's define our model and verify that it can process a single input/output.\n",
    "\n",
    "For this we'll use Thinc and the `define_operators` context manager to connect the pieces together using overloaded operators for `chain` and `clone` operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from thinc.model import Model\n",
    "from thinc.api import concatenate, chain, clone, list2ragged\n",
    "from thinc.api import reduce_sum, Mish, with_array, Embed, residual\n",
    "\n",
    "def build_model(n_hidden: int, dropout: float = 0.1) -> ModelT:\n",
    "    with Model.define_operators({\">>\": chain, \"|\": concatenate, \"**\": clone}):\n",
    "        model = (\n",
    "            # Iterate over each element in the batch\n",
    "            with_array(\n",
    "                # Embed the vocab indices\n",
    "                Embed(n_hidden, len(vocab), column=0)\n",
    "                # Activate each batch of embedding sequences separately first\n",
    "                >> Mish(n_hidden, dropout=dropout)\n",
    "            )\n",
    "            # Convert to ragged so we can use the reduction layers\n",
    "            >> list2ragged()\n",
    "            # Sum the features for each batch input\n",
    "            >> reduce_sum()\n",
    "            # Process with a small resnet\n",
    "            >> residual(Mish(n_hidden, normalize=True)) ** 4\n",
    "            # Convert (batch_size, n_hidden) to (batch_size, 1)\n",
    "            >> Mish(1)\n",
    "        )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It\n",
    "\n",
    "Let's pass an example through the model to make sure we have all the sizes right."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text, X, Y = to_example(\"14x + 2y - 3x + 7x\")\n",
    "m = build_model(12)\n",
    "m.initialize([X], m.ops.asarray(Y, dtype=\"f\"))\n",
    "mY = m.predict([X])\n",
    "print(mY.shape)\n",
    "assert mY.shape == (1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Training Datasets\n",
    "\n",
    "Now that we can generate examples and we have a model that can process them, let's generate random unique training and evaluation datasets.\n",
    "\n",
    "For this we'll write another helper function that can generate (n) training examples and respects an exclude list to avoid letting examples from the training/test sets overlap."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple, Optional, Set, List\n",
    "\n",
    "DatasetTuple = Tuple[List[str], List[ModelX], List[ModelY]]\n",
    "\n",
    "def generate_dataset(\n",
    "    size: int,\n",
    "    exclude: Optional[Set[str]] = None,\n",
    ") -> DatasetTuple:\n",
    "    ops: Ops = get_current_ops()\n",
    "    texts: List[str] = generate_problems(size, exclude=exclude)\n",
    "    examples: List[ModelX] = []\n",
    "    labels: List[ModelY] = []\n",
    "    for i, text in enumerate(texts):\n",
    "        text, x, y = to_example(text)\n",
    "        examples.append(x)\n",
    "        labels.append(y)\n",
    "\n",
    "    return texts, examples, labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It\n",
    "\n",
    "Generate a small dataset to be sure everything is working as expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts, x, y = generate_dataset(10)\n",
    "assert len(texts) == 10\n",
    "assert len(x) == 10\n",
    "assert len(y) == 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate Model Performance\n",
    "\n",
    "We're almost ready to train our model, we just need to write a function that will check a given trained model against a given dataset and return a 0-1 score of how accurate it was.\n",
    "\n",
    "We'll use this function to print the score as training progresses and print final test predictions at the end of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from wasabi import msg\n",
    "\n",
    "def evaluate_model(\n",
    "    model: ModelT,\n",
    "    *,\n",
    "    print_problems: bool = False,\n",
    "    texts: List[str],\n",
    "    X: List[ModelX],\n",
    "    Y: List[ModelY],\n",
    "):\n",
    "    Yeval = model.predict(X)\n",
    "    correct_count = 0\n",
    "    print_n = 12\n",
    "    if print_problems:\n",
    "        msg.divider(f\"eval samples max({print_n})\")\n",
    "    for text, y_answer, y_guess in zip(texts, Y, Yeval):\n",
    "        y_guess = round(float(y_guess))\n",
    "        correct = y_guess == int(y_answer)\n",
    "        print_fn = msg.fail\n",
    "        if correct:\n",
    "            correct_count += 1\n",
    "            print_fn = msg.good\n",
    "        if print_problems and print_n > 0:\n",
    "            print_n -= 1\n",
    "            print_fn(f\"Answer[{int(y_answer[0])}] Guess[{y_guess}] Text: {text}\")\n",
    "    if print_problems:\n",
    "        print(f\"Model predicted {correct_count} out of {len(X)} correctly.\")\n",
    "    return correct_count / len(X)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try It\n",
    "\n",
    "Let's try it out with an untrained model and expect to see a really sad score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "texts, X, Y = generate_dataset(128)\n",
    "m = build_model(12)\n",
    "m.initialize(X, m.ops.asarray(Y, dtype=\"f\"))\n",
    "# Assume the model should do so poorly as to round down to 0\n",
    "assert round(evaluate_model(m, texts=texts, X=X, Y=Y)) == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train/Evaluate a Model\n",
    "\n",
    "The final helper function we need is one to train and evaluate a model given two input datasets. \n",
    "\n",
    "This function does a few things:\n",
    "\n",
    " 1. Create an Adam optimizer we can use for minimizing the model's prediction error.\n",
    " 2. Loop over the given training dataset (epoch) number of times.\n",
    " 3. For each epoch, make batches of (batch_size) examples. For each batch(X), predict the number of like terms (Yh) and subtract the known answers (Y) to get the prediction error. Update the model using the optimizer with the calculated error.\n",
    " 5. After each epoch, check the model performance against the evaluation dataset.\n",
    " 6. Save the model weights for the best score out of all the training epochs.\n",
    " 7. After all training is done, restore the best model and print results from the evaluation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from thinc.api import Adam\n",
    "from wasabi import msg\n",
    "import numpy\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "def train_and_evaluate(\n",
    "    model: ModelT,\n",
    "    train_tuple: DatasetTuple,\n",
    "    eval_tuple: DatasetTuple,\n",
    "    *,\n",
    "    lr: float = 3e-3,\n",
    "    batch_size: int = 64,\n",
    "    epochs: int = 48,\n",
    ") -> float:\n",
    "    (train_texts, train_X, train_y) = train_tuple\n",
    "    (eval_texts, eval_X, eval_y) = eval_tuple\n",
    "    msg.divider(\"Train and Evaluate Model\")\n",
    "    msg.info(f\"Batch size = {batch_size}\\tEpochs = {epochs}\\tLearning Rate = {lr}\")\n",
    "\n",
    "    optimizer = Adam(lr)\n",
    "    best_score: float = 0.0\n",
    "    best_model: Optional[bytes] = None\n",
    "    for n in range(epochs):\n",
    "        loss = 0.0\n",
    "        batches = model.ops.multibatch(batch_size, train_X, train_y, shuffle=True)\n",
    "        for X, Y in tqdm(batches, leave=False, unit=\"batches\"):\n",
    "            Y = model.ops.asarray(Y, dtype=\"float32\")\n",
    "            Yh, backprop = model.begin_update(X)\n",
    "            err = Yh - Y\n",
    "            backprop(err)\n",
    "            loss += (err ** 2).sum()\n",
    "            model.finish_update(optimizer)\n",
    "        score = evaluate_model(model, texts=eval_texts, X=eval_X, Y=eval_y)\n",
    "        if score > best_score:\n",
    "            best_model = model.to_bytes()\n",
    "            best_score = score\n",
    "        print(f\"{n}\\t{score:.2f}\\t{loss:.2f}\")\n",
    "\n",
    "    if best_model is not None:\n",
    "        model.from_bytes(best_model)\n",
    "    print(f\"Evaluating with best model\")\n",
    "    score = evaluate_model(\n",
    "        model, texts=eval_texts, print_problems=True, X=eval_X, Y=eval_y\n",
    "    )\n",
    "    print(f\"Final Score: {score}\")\n",
    "    return score\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll generate the dataset first, so we can iterate on the model without having to spend time generating examples for each run. This also ensures we have the same dataset across different model runs, to make it easier to compare performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_size = 1024 * 8\n",
    "test_size = 2048\n",
    "seen_texts: Set[str] = set()\n",
    "with msg.loading(f\"Generating train dataset with {train_size} examples...\"):\n",
    "    train_dataset = generate_dataset(train_size, seen_texts)\n",
    "msg.good(f\"Train set created with {train_size} examples.\")\n",
    "with msg.loading(f\"Generating eval dataset with {test_size} examples...\"):\n",
    "    eval_dataset = generate_dataset(test_size, seen_texts)\n",
    "msg.good(f\"Eval set created with {test_size} examples.\")\n",
    "init_x = train_dataset[1][:2]\n",
    "init_y = train_dataset[2][:2]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can build, train, and evaluate our model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model(64)\n",
    "model.initialize(init_x, init_y)\n",
    "train_and_evaluate(\n",
    "    model, train_dataset, eval_dataset, lr=2e-3, batch_size=64, epochs=16\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Intermediate Exercise\n",
    "\n",
    "The model we built can train up to ~80% given 100 or more epochs. Improve the model architecture so that it trains to a similar accuracy while requiring fewer epochs or a smaller dataset size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from thinc.model import Model\n",
    "from thinc.types import Array2d, Array1d\n",
    "from thinc.api import chain, clone, list2ragged, reduce_mean, Mish, with_array, Embed, residual\n",
    "\n",
    "def custom_model(n_hidden: int, dropout: float = 0.1) -> Model[List[Array2d], Array2d]:\n",
    "    # Put your custom architecture here\n",
    "    return build_model(n_hidden, dropout)\n",
    "\n",
    "model = custom_model(64)\n",
    "model.initialize(init_x, init_y)\n",
    "train_and_evaluate(\n",
    "    model, train_dataset, eval_dataset, lr=2e-3, batch_size=64, epochs=16\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Advanced Exercise\n",
    "\n",
    "Rewrite the model to encode the whole expression with a BiLSTM, and then generate pairs of terms, using the BiLSTM vectors. Over each pair of terms, predict whether the terms are alike or unlike."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "from thinc.types import Array2d, Ragged\n",
    "from thinc.model import Model\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Comparisons:\n",
    "    data: Array2d  # Batch of vectors for each pair\n",
    "    indices: Array2d  # Int array of shape (N, 3), showing the (batch, term1, term2) positions\n",
    "\n",
    "def pairify() -> Model[Ragged, Comparisons]:\n",
    "    \"\"\"Create pair-wise comparisons for items in a sequence. For each sequence of N\n",
    "    items, there will be (N**2-N)/2 comparisons.\"\"\"\n",
    "    ...\n",
    "\n",
    "def predict_over_pairs(model: Model[Array2d, Array2d]) -> Model[Comparisons, Comparisons]:\n",
    "    \"\"\"Apply a prediction model over a batch of comparisons. Outputs a Comparisons\n",
    "    object where the data is the scores. The prediction model should predict over\n",
    "    two classes, True and False.\"\"\"\n",
    "    ...\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}