vmal/3-digit-basic-calc

1{2 "cells": [3  {4   "cell_type": "markdown",5   "metadata": {},6   "source": [7    "# \ud83e\uddee 3-digit-basic-calc \u2014 3-digit arithmetic with a 1.6M-parameter transformer\n",8    "\n",9    "A from-scratch transformer that does `+ \u2212 \u00d7 \u00f7` on 3-digit numbers by writing out\n",10    "the algorithm step by step. This notebook loads the model and runs it.\n",11    "\n",12    "Model: [`vmal/3-digit-basic-calc`](https://huggingface.co/vmal/3-digit-basic-calc)\n"13   ]14  },15  {16   "cell_type": "code",17   "metadata": {},18   "execution_count": null,19   "outputs": [],20   "source": [21    "%pip install -q transformers==5.3.0 torch safetensors\n"22   ]23  },24  {25   "cell_type": "markdown",26   "metadata": {},27   "source": [28    "## Load the model\n"29   ]30  },31  {32   "cell_type": "code",33   "metadata": {},34   "execution_count": null,35   "outputs": [],36   "source": [37    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",38    "\n",39    "REPO = \"vmal/3-digit-basic-calc\"\n",40    "model = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True).eval()\n",41    "tok   = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)\n",42    "print(\"parameters:\", sum(p.numel() for p in model.parameters()))\n"43   ]44  },45  {46   "cell_type": "markdown",47   "metadata": {},48   "source": [49    "## Solve some problems\n",50    "\n",51    "`model.solve(tokenizer, expression)` returns the human-readable answer.\n"52   ]53  },54  {55   "cell_type": "code",56   "metadata": {},57   "execution_count": null,58   "outputs": [],59   "source": [60    "for e in [\"842/37\", \"213*145\", \"999*999\", \"-500+500\", \"3/31\", \"12/0\"]:\n",61    "    print(f\"{e:>10} = {model.solve(tok, e)}\")\n"62   ]63  },64  {65   "cell_type": "markdown",66   "metadata": {},67   "source": [68    "## See the model's reasoning\n",69    "\n",70    "Pass `return_trace=True` to get the raw scratchpad the model generates \u2014\n",71    "each `<step>`/`<qmul>`/`<rem>` is one local computation.\n"72   ]73  },74  {75   "cell_type": "code",76   "metadata": {},77   "execution_count": null,78   "outputs": [],79   "source": [80    "answer, trace = model.solve(tok, \"842/37\", return_trace=True)\n",81    "print(\"answer:\", answer)\n",82    "print()\n",83    "for c in ['<div>','<mul>','<add>','<sub>','<state>','<step>','<qmul>','<rem>','<col>','<ans>']:\n",84    "    trace = trace.replace(c, '\\n'+c+' ')\n",85    "print(trace.strip())\n"86   ]87  },88  {89   "cell_type": "markdown",90   "metadata": {},91   "source": [92    "## Quick in-range behavior check on random problems\n",93    "\n",94    "This small smoke test samples the supported operand range. It does not load\n",95    "the training prompts, so it should not be described as a leakage-controlled\n",96    "unseen benchmark; the card reports that benchmark separately.\n"97   ]98  },99  {100   "cell_type": "code",101   "metadata": {},102   "execution_count": null,103   "outputs": [],104   "source": [105    "import random\n",106    "random.seed(0)\n",107    "def truth(a, b, op):\n",108    "    if op == '+': return str(a + b)\n",109    "    if op == '-': return str(a - b)\n",110    "    if op == '*': return str(a * b)\n",111    "    if b == 0: return 'NAN'\n",112    "    # Exact integer round-half-up to three decimals (no float/banker's rounding).\n",113    "    denominator = abs(b)\n",114    "    scaled, remainder = divmod(abs(a) * 1000, denominator)\n",115    "    scaled += int(2 * remainder >= denominator)\n",116    "    integer, fraction = divmod(scaled, 1000)\n",117    "    magnitude = str(integer)\n",118    "    if fraction:\n",119    "        magnitude += '.' + f'{fraction:03d}'.rstrip('0')\n",120    "    negative = (a < 0) != (b < 0)\n",121    "    return '-' + magnitude if negative and scaled else magnitude\n",122    "good=n=0\n",123    "for _ in range(30):\n",124    "    a=random.randint(-999,999); b=random.randint(-999,999); op=random.choice('+-*/')\n",125    "    if op=='/' and b==0: b=7\n",126    "    got=model.solve(tok, f'{a}{op}{b}'); exp=truth(a,b,op)\n",127    "    good+= (got==exp); n+=1\n",128    "print(f'{good}/{n} correct on random in-range problems')\n"129   ]130  }131 ],132 "metadata": {133  "kernelspec": {134   "display_name": "Python 3",135   "language": "python",136   "name": "python3"137  },138  "language_info": {139   "name": "python"140  }141 },142 "nbformat": 4,143 "nbformat_minor": 5144}145