Duplicate from cfahlgren1/enron-autocomplete

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+hero.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,111 @@
+---
+library_name: onnx
+pipeline_tag: text-generation
+datasets:
+  - corbt/enron-emails
+language:
+  - en
+tags:
+  - autocomplete
+  - gru
+  - onnxruntime-web
+  - in-browser
+---
+
+# enron-autocomplete
+
+[![Typing "Please accept my resignation. I will be o" and the model suggesting "ut of the office"](hero.png)](https://huggingface.co/spaces/cfahlgren1/enron-autocomplete)
+
+A tiny email autocomplete that runs in your browser. It learned to write from Enron's emails.
+
+**[Try the demo](https://huggingface.co/spaces/cfahlgren1/enron-autocomplete)**
+
+| | |
+|---|---|
+| Size | 944K parameters, 0.98 MB |
+| Speed | about 4 ms per suggestion, on CPU in Chrome |
+| Training | 14 minutes on one A10G with Hugging Face Jobs, 35 cents |
+| Data | 119,859 Enron emails |
+
+## What it does
+
+It finishes the word or short phrase you're typing.
+
+| You type | It suggests |
+|---|---|
+| Please accept my resig | nation |
+| Please accept my resignation. I will be o | ut of the office |
+| The books are fine. Please call me i | f you have any |
+| Enron N | orth America Corp. |
+
+## How sure it has to be
+
+The model knows how sure it is of each guess. You pick a cutoff, and it only suggests when it's at least that sure. A higher cutoff means fewer suggestions, and more of them right.
+
+| Cutoff | How often it suggests | How often it's right |
+|---|---|---|
+| 20% | 30% of the time | 48% |
+| 35% | 18% of the time | 66% |
+| 50% | 11% of the time | 80% |
+| 70% | 7% of the time | 90% |
+
+The demo uses 40%.
+
+These numbers come from 5,000 places in emails from people the model never saw in training. "Right" means the suggestion exactly matched what the person typed next. Right suggestions are about 6 characters long on average.
+
+## Files
+
+| File | What it is |
+|---|---|
+| `model_quint8.onnx` | The model to use. 0.98 MB. |
+| `model.onnx` | The same model before shrinking. 3.8 MB. |
+| `best.pt` | The PyTorch checkpoint. |
+| `vocab.json` | The character ids. |
+| `report.json` | The full training and test results. |
+| `train.py` | The training script. |
+
+I also tried a bigger version with 3.46M parameters and 3.5 MB. It wasn't better. When both models suggest equally often, 9% of the time, the small one is right about 85% of the time and the big one 82%. The big one got the same training settings and may need different ones to do well.
+
+## Use it
+
+The model takes one character and its memory of the text so far. It returns its guess for the next character and an updated memory. So each key press is one small step.
+
+```js
+import * as ort from "onnxruntime-web";
+
+const session = await ort.InferenceSession.create("model_quint8.onnx", {
+  executionProviders: ["wasm"],
+});
+let state = new ort.Tensor("float32", new Float32Array(512), [1, 512]);
+
+async function step(id) {
+  const input_id = new ort.Tensor("int64", BigInt64Array.of(BigInt(id)), [1]);
+  const { logits, next_state } = await session.run({ input_id, state });
+  state = next_state;
+  return logits.data;
+}
+
+let logits = await step(1); // 1 starts a new email
+for (const char of "Let me kn") logits = await step(char.charCodeAt(0) - 29);
+```
+
+A character's id is its character code minus 29. A new line is 2 and the end of an email is 1. Only English keyboard characters are supported.
+
+To build a suggestion, keep taking the most likely next character and multiply the probabilities together as you go. Stop at a new line, or when that running product drops below your cutoff. Show the text up to the last finished word.
+
+## Training
+
+It read 40 million characters from 119,859 emails sent from enron.com addresses, 30 times over. Typos were added to what it read, so it keeps predicting when you make one. The run took 14 minutes on one A10G with Hugging Face Jobs and cost 35 cents.
+
+[`train.py`](train.py) is the whole thing: it downloads the emails, trains, tests and exports the model. To train your own:
+
+```bash
+hf jobs uv run --flavor a10g-large --secrets HF_TOKEN train.py --repo your-name/your-model
+```
+
+## Limits
+
+- It only knows office email from 2001. Outside that it has little to say.
+- With no cutoff it rambles. "Dear Santa," becomes "I am attaching a copy of the agreement to the Enron Corp."
+- It learned real names and sign-offs from the emails, so keep suggestions short.
+- English only.

best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f2a663ad7a697a00ebf6bd15942824fcefa46975f9935c7ac6dd1eceedf3a36
+size 3780287

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ccc402ffaa27e6382114ad0d44a2e5a01e9beabe9d81e3c3ae0ad2f9676627b
+size 3780016

model_quint8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:079424a9d2a72b9e5b2d6002c6466b836023c088780b7d292748f37564dc143b
+size 981294

report.json

@@ -0,0 +1,548 @@
+{
+  "args": {
+    "hidden": 512,
+    "epochs": 30,
+    "batch": 256,
+    "seq": 256,
+    "burn_in": 32,
+    "lr": 0.0015,
+    "typo_rate": 0.08,
+    "limit_frac": 1.0,
+    "eval_positions": 5000,
+    "context": 128,
+    "horizon": 48,
+    "repo": "cfahlgren1/email-autocomplete-gru",
+    "run_name": "gru-512",
+    "no_push": false
+  },
+  "device": "cuda",
+  "parameters": 944354,
+  "data": {
+    "train": {
+      "emails": 119859,
+      "chars": 40136658
+    },
+    "val": {
+      "emails": 8252,
+      "chars": 2775980
+    },
+    "test": {
+      "emails": 6922,
+      "chars": 2710169
+    }
+  },
+  "history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.7093,
+      "val_loss": 1.3716,
+      "seconds": 23.9,
+      "chars_per_sec": 1678766
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.2792,
+      "val_loss": 1.2911,
+      "seconds": 23.6,
+      "chars_per_sec": 1701572
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.2231,
+      "val_loss": 1.2582,
+      "seconds": 23.6,
+      "chars_per_sec": 1700904
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.195,
+      "val_loss": 1.2396,
+      "seconds": 23.6,
+      "chars_per_sec": 1702926
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.1768,
+      "val_loss": 1.2254,
+      "seconds": 23.5,
+      "chars_per_sec": 1703932
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.1639,
+      "val_loss": 1.2171,
+      "seconds": 23.6,
+      "chars_per_sec": 1701905
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.1538,
+      "val_loss": 1.2087,
+      "seconds": 23.6,
+      "chars_per_sec": 1701450
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.1456,
+      "val_loss": 1.2043,
+      "seconds": 23.6,
+      "chars_per_sec": 1702601
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.1388,
+      "val_loss": 1.1969,
+      "seconds": 23.5,
+      "chars_per_sec": 1703175
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.1327,
+      "val_loss": 1.193,
+      "seconds": 23.6,
+      "chars_per_sec": 1701059
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.1276,
+      "val_loss": 1.1895,
+      "seconds": 23.6,
+      "chars_per_sec": 1702827
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.1228,
+      "val_loss": 1.1856,
+      "seconds": 23.6,
+      "chars_per_sec": 1700650
+    },
+    {
+      "epoch": 13,
+      "train_loss": 1.1186,
+      "val_loss": 1.1824,
+      "seconds": 23.5,
+      "chars_per_sec": 1703353
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.1141,
+      "val_loss": 1.1791,
+      "seconds": 23.6,
+      "chars_per_sec": 1702586
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.1108,
+      "val_loss": 1.1765,
+      "seconds": 23.6,
+      "chars_per_sec": 1702570
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.1075,
+      "val_loss": 1.175,
+      "seconds": 23.6,
+      "chars_per_sec": 1702686
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.1041,
+      "val_loss": 1.1718,
+      "seconds": 23.5,
+      "chars_per_sec": 1703957
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.1015,
+      "val_loss": 1.17,
+      "seconds": 23.6,
+      "chars_per_sec": 1701778
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.0985,
+      "val_loss": 1.1678,
+      "seconds": 23.6,
+      "chars_per_sec": 1702612
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.0959,
+      "val_loss": 1.1659,
+      "seconds": 23.6,
+      "chars_per_sec": 1702960
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.0935,
+      "val_loss": 1.1642,
+      "seconds": 23.6,
+      "chars_per_sec": 1701776
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.0913,
+      "val_loss": 1.1632,
+      "seconds": 23.6,
+      "chars_per_sec": 1702317
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.0896,
+      "val_loss": 1.1624,
+      "seconds": 23.6,
+      "chars_per_sec": 1701215
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.0874,
+      "val_loss": 1.1613,
+      "seconds": 23.6,
+      "chars_per_sec": 1701278
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.0863,
+      "val_loss": 1.1606,
+      "seconds": 23.6,
+      "chars_per_sec": 1700749
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.085,
+      "val_loss": 1.1596,
+      "seconds": 23.6,
+      "chars_per_sec": 1701575
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.0839,
+      "val_loss": 1.1592,
+      "seconds": 23.6,
+      "chars_per_sec": 1699705
+    },
+    {
+      "epoch": 28,
+      "train_loss": 1.083,
+      "val_loss": 1.1588,
+      "seconds": 23.6,
+      "chars_per_sec": 1700990
+    },
+    {
+      "epoch": 29,
+      "train_loss": 1.0823,
+      "val_loss": 1.1588,
+      "seconds": 23.6,
+      "chars_per_sec": 1702216
+    },
+    {
+      "epoch": 30,
+      "train_loss": 1.0821,
+      "val_loss": 1.1586,
+      "seconds": 23.5,
+      "chars_per_sec": 1704907
+    }
+  ],
+  "best_val_loss": 1.1586,
+  "test_positions": 5000,
+  "held_out_sender_suggestions": {
+    "0.05": {
+      "coverage": 0.624,
+      "precision": 0.2064,
+      "avg_chars_when_correct": 6.93,
+      "chars_saved_per_position": 0.892
+    },
+    "0.1": {
+      "coverage": 0.4622,
+      "precision": 0.3124,
+      "avg_chars_when_correct": 6.41,
+      "chars_saved_per_position": 0.925
+    },
+    "0.2": {
+      "coverage": 0.2974,
+      "precision": 0.4808,
+      "avg_chars_when_correct": 6.03,
+      "chars_saved_per_position": 0.863
+    },
+    "0.35": {
+      "coverage": 0.1796,
+      "precision": 0.6559,
+      "avg_chars_when_correct": 5.8,
+      "chars_saved_per_position": 0.683
+    },
+    "0.5": {
+      "coverage": 0.1114,
+      "precision": 0.7953,
+      "avg_chars_when_correct": 5.64,
+      "chars_saved_per_position": 0.5
+    },
+    "0.7": {
+      "coverage": 0.0694,
+      "precision": 0.8991,
+      "avg_chars_when_correct": 5.4,
+      "chars_saved_per_position": 0.337
+    }
+  },
+  "exports": {
+    "float": {
+      "size_mb": 3.78,
+      "native_cpu_single_thread_step": {
+        "p50_ms": 0.114,
+        "p95_ms": 0.125
+      },
+      "suggestions_1000_positions": {
+        "0.05": {
+          "coverage": 0.638,
+          "precision": 0.2132,
+          "avg_chars_when_correct": 6.7,
+          "chars_saved_per_position": 0.911
+        },
+        "0.1": {
+          "coverage": 0.496,
+          "precision": 0.3085,
+          "avg_chars_when_correct": 6.3,
+          "chars_saved_per_position": 0.964
+        },
+        "0.2": {
+          "coverage": 0.312,
+          "precision": 0.4583,
+          "avg_chars_when_correct": 5.94,
+          "chars_saved_per_position": 0.849
+        },
+        "0.35": {
+          "coverage": 0.202,
+          "precision": 0.604,
+          "avg_chars_when_correct": 5.77,
+          "chars_saved_per_position": 0.704
+        },
+        "0.5": {
+          "coverage": 0.119,
+          "precision": 0.7479,
+          "avg_chars_when_correct": 5.58,
+          "chars_saved_per_position": 0.497
+        },
+        "0.7": {
+          "coverage": 0.078,
+          "precision": 0.8718,
+          "avg_chars_when_correct": 5.21,
+          "chars_saved_per_position": 0.354
+        }
+      }
+    },
+    "quint8": {
+      "size_mb": 0.98,
+      "native_cpu_single_thread_step": {
+        "p50_ms": 0.059,
+        "p95_ms": 0.084
+      },
+      "suggestions_1000_positions": {
+        "0.05": {
+          "coverage": 0.655,
+          "precision": 0.2,
+          "avg_chars_when_correct": 6.83,
+          "chars_saved_per_position": 0.895
+        },
+        "0.1": {
+          "coverage": 0.492,
+          "precision": 0.311,
+          "avg_chars_when_correct": 6.35,
+          "chars_saved_per_position": 0.971
+        },
+        "0.2": {
+          "coverage": 0.318,
+          "precision": 0.4686,
+          "avg_chars_when_correct": 5.97,
+          "chars_saved_per_position": 0.89
+        },
+        "0.35": {
+          "coverage": 0.198,
+          "precision": 0.6414,
+          "avg_chars_when_correct": 5.66,
+          "chars_saved_per_position": 0.719
+        },
+        "0.5": {
+          "coverage": 0.123,
+          "precision": 0.7642,
+          "avg_chars_when_correct": 5.53,
+          "chars_saved_per_position": 0.52
+        },
+        "0.7": {
+          "coverage": 0.081,
+          "precision": 0.8889,
+          "avg_chars_when_correct": 5.42,
+          "chars_saved_per_position": 0.39
+        }
+      }
+    }
+  },
+  "prompt_samples": {
+    "Hi John,\n\nThanks for ": {
+      "continuation": "the info. I will be out of the office tomorrow a",
+      "suggestion_at_0.2": "the"
+    },
+    "Please let me know if ": {
+      "continuation": "you have any questions.\n\nThanks,\n\nKayThe follow",
+      "suggestion_at_0.2": "you have any"
+    },
+    "I will send you ": {
+      "continuation": "a copy of the attached document. I will be out o",
+      "suggestion_at_0.2": "a"
+    },
+    "Let me kn": {
+      "continuation": "ow if you have any questions.\n\nThanks,\n\nKayThe ",
+      "suggestion_at_0.2": "ow if you have"
+    },
+    "Thnaks for the ": {
+      "continuation": "info. I will be out of the office tomorrow and w",
+      "suggestion_at_0.2": ""
+    },
+    "I look forward to ": {
+      "continuation": "seeing you there. I will be out of the office to",
+      "suggestion_at_0.2": ""
+    },
+    "Can we schedule a ": {
+      "continuation": "meeting to discuss the attached spreadsheet for ",
+      "suggestion_at_0.2": "meeting"
+    },
+    "Sorry for the dela": {
+      "continuation": "y. I will be out of the office tomorrow and we w",
+      "suggestion_at_0.2": "y."
+    },
+    "Plese find attached ": {
+      "continuation": "the following contract for the Enron Corp. Gas P",
+      "suggestion_at_0.2": "the"
+    },
+    "We need to discuss ": {
+      "continuation": "the attached spreadsheet for the following dates",
+      "suggestion_at_0.2": "the"
+    }
+  },
+  "held_out_examples": [
+    {
+      "context": "ou trading a book yet? Playing any poker? Bingham was here f",
+      "truth": "or Q-School last week. Didn't go so well.",
+      "suggestion_at_0.2": "or"
+    },
+    {
+      "context": "Please use this model. I had a few change",
+      "truth": "s after talking to the guys in the west. Thanks!",
+      "suggestion_at_0.2": "s to"
+    },
+    {
+      "context": "not open the attachment. Please, send it as a MS Word docume",
+      "truth": "nt or a .pdf file.",
+      "suggestion_at_0.2": "nt."
+    },
+    {
+      "context": "n I do to help. I would also like to visit on status of e*p ",
+      "truth": "group and mariner.",
+      "suggestion_at_0.2": ""
+    },
+    {
+      "context": " complete hourly schedule for June 1, including IBM and inco",
+      "truth": "rporating all known changes confirmed by that ti",
+      "suggestion_at_0.2": "me"
+    },
+    {
+      "context": " delivery rights at SoCal/Topock for at least a pro rata por",
+      "truth": "tion of its released capacity.",
+      "suggestion_at_0.2": "tion of"
+    },
+    {
+      "context": "i talked to chris a li",
+      "truth": "ttle bit and he is going to call back tonight at",
+      "suggestion_at_0.2": "ttle"
+    },
+    {
+      "context": " I might be missing something about the way you guys do thes",
+      "truth": "e deals in Canada so I thought I would run this ",
+      "suggestion_at_0.2": "e"
+    },
+    {
+      "context": "Traci,\n\nThanks. This is the best solution given all the unce",
+      "truth": "rtainties we face.",
+      "suggestion_at_0.2": "rtainty"
+    },
+    {
+      "context": "\nAquila,s strategy comes as little surprise to Stuebi. As mu",
+      "truth": "ch success as the company has had, he said he th",
+      "suggestion_at_0.2": "ch"
+    },
+    {
+      "context": "Jim,\n\nWhat about 2 p.m., my time, on Friday, Au",
+      "truth": "gust 17. If this works for you, I shall try to i",
+      "suggestion_at_0.2": "gust"
+    },
+    {
+      "context": "'s and CDD's for Philedelphia, Allentown Washington DC, Balt",
+      "truth": "imore, Pittsburg, Newark, New York City, Cinncin",
+      "suggestion_at_0.2": "imore"
+    },
+    {
+      "context": " release.\n\nThe company will offer derivatives to hedge volat",
+      "truth": "ile electricity prices.",
+      "suggestion_at_0.2": "ility"
+    },
+    {
+      "context": "day, I need access to West deals in Sitara. At a mininum I s",
+      "truth": "till need access to all Chicago/Midwest deals al",
+      "suggestion_at_0.2": "ent"
+    },
+    {
+      "context": "on, attached for your review is a draft of the First Amendme",
+      "truth": "nt to Master Agreement between Enron North Ameri",
+      "suggestion_at_0.2": "nt to"
+    },
+    {
+      "context": "I have provided all the info that Tami Scudder ",
+      "truth": "requested today.",
+      "suggestion_at_0.2": "is"
+    },
+    {
+      "context": "ontana Power Company (MPC) System Border for Thursday and Fr",
+      "truth": "iday (July 12, 13) @ $67.85/MWh. Thanks, John",
+      "suggestion_at_0.2": "iday"
+    },
+    {
+      "context": "e EnBank contracts physical or financial? It seems to me we ",
+      "truth": "should use the same customer profiles - that is ",
+      "suggestion_at_0.2": ""
+    },
+    {
+      "context": "I have an appointment at 1:00 today and should be ",
+      "truth": "back to the office between 2 - 2:30.",
+      "suggestion_at_0.2": ""
+    },
+    {
+      "context": "I'm glad everythi",
+      "truth": "ng worked out for you. Hope you have a great tim",
+      "suggestion_at_0.2": "ng is"
+    },
+    {
+      "context": "o's Line 1100 from Eunice Station to Pecos River Station wil",
+      "truth": "l be out of service for 16 hours Saturday beginn",
+      "suggestion_at_0.2": "l be"
+    },
+    {
+      "context": "rty\" in the second line of paragraph 1. Otherwise they look ",
+      "truth": "great to me.",
+      "suggestion_at_0.2": ""
+    },
+    {
+      "context": "on Houston in the next day or so.\n\nPlease call me if you hav",
+      "truth": "e any questions.",
+      "suggestion_at_0.2": "e any questions."
+    },
+    {
+      "context": "s not guarantee any continuity of support for you.\n\nPlease, ",
+      "truth": "let me know what you decide to do over the next ",
+      "suggestion_at_0.2": ""
+    },
+    {
+      "context": " tend to agree with you, but Nora Dobin is the lawyer who kn",
+      "truth": "ows the most about the DealBench system - you sh",
+      "suggestion_at_0.2": "ows"
+    }
+  ],
+  "total_seconds": 831.5
+}

train.py

@@ -0,0 +1,557 @@
+# /// script
+# requires-python = ">=3.11,<3.14"
+# dependencies = [
+#   "torch>=2.6,<2.8",
+#   "datasets>=3.0",
+#   "numpy<3",
+#   "onnx>=1.17",
+#   "onnxruntime>=1.20",
+#   "huggingface_hub>=0.30",
+# ]
+# ///
+"""Train a character-level GRU for in-browser email autocomplete on Enron."""
+
+import argparse
+import hashlib
+import json
+import os
+import re
+import time
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from datasets import load_dataset
+from huggingface_hub import HfApi
+from onnxruntime.quantization import QuantType, quantize_dynamic
+
+# Vocabulary: PAD, EOS, newline, then printable ASCII.
+PAD, EOS, NEWLINE, FIRST_PRINTABLE = 0, 1, 2, 3
+VOCAB_SIZE = FIRST_PRINTABLE + (127 - 32)
+IGNORE = -100
+
+ENCODE = np.full(256, 255, dtype=np.uint8)
+ENCODE[10] = NEWLINE
+ENCODE[32:127] = np.arange(FIRST_PRINTABLE, VOCAB_SIZE, dtype=np.uint8)
+DECODE = {NEWLINE: "\n", **{FIRST_PRINTABLE + i: chr(32 + i) for i in range(95)}}
+
+LETTER = np.zeros(VOCAB_SIZE, dtype=bool)
+for char in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ":
+    LETTER[ENCODE[ord(char)]] = True
+
+KEYBOARD_NEIGHBORS = {
+    "q": "wa", "w": "qes", "e": "wrd", "r": "etf", "t": "ryg", "y": "tuh", "u": "yij",
+    "i": "uok", "o": "ipl", "p": "ol", "a": "qsz", "s": "awdx", "d": "serf", "f": "drgv",
+    "g": "fthb", "h": "gyjn", "j": "hukm", "k": "jil", "l": "kop", "z": "asx", "x": "zsc",
+    "c": "xdv", "v": "cfb", "b": "vgn", "n": "bhm", "m": "njk",
+}
+NEIGHBOR_TABLE = np.tile(np.arange(VOCAB_SIZE, dtype=np.uint8)[:, None], (1, 4))
+for key, neighbors in KEYBOARD_NEIGHBORS.items():
+    for slot in range(4):
+        NEIGHBOR_TABLE[ENCODE[ord(key)], slot] = ENCODE[ord(neighbors[slot % len(neighbors)])]
+
+CUT_REPLY = re.compile(
+    r"(?im)^\s*(?:-+\s*original message|-+\s*forwarded by|-+\s*inline attachment|from:\s|to:\s|sent:\s"
+    r"|subject:\s|cc:\s|.*\bwrote:\s*$|_{5,}|-{5,}|={5,}|\*{5,}|\d{1,2}/\d{1,2}/\d{2,4}\s+\d{1,2}:\d{2})"
+)
+SENSITIVE = re.compile(
+    r"(?i)https?://|www\.|\S+@\S+|\(?\b\d{3}\)?[-. ]\d{3}[-. ]\d{4}\b|\b\d{3}-\d{2}-\d{4}\b"
+)
+BULK_SENDER = re.compile(r"(?i)announce|newsletter|mailer|outlook\.team|^all[._]|general\.|chairman|enron\.(?:on|net)")
+SMART_CHARS = str.maketrans({"‘": "'", "’": "'", "“": '"', "”": '"', "–": "-", "—": "-", "\t": " "})
+
+
+def clean_paragraphs(body: str) -> list[str]:
+    text = body.replace("\r\n", "\n").replace("=20", " ").replace("=\n", "").translate(SMART_CHARS)
+    text = "\n".join(line for line in text.split("\n") if not line.lstrip().startswith(">"))
+    cut = CUT_REPLY.search(text)
+    if cut:
+        text = text[: cut.start()]
+    text = re.sub(r"(?m)^-- *\n[\s\S]*", "", text)
+
+    paragraphs = []
+    for block in re.split(r"\n\s*\n", text):
+        lines = [re.sub(r" +", " ", line).strip() for line in block.split("\n") if line.strip()]
+        if not lines:
+            continue
+        # Enron bodies are hard-wrapped near 70 columns; only keep breaks after short lines.
+        joined = lines[0]
+        for previous, line in zip(lines, lines[1:]):
+            joined += (" " if len(previous) >= 50 else "\n") + line
+        letters = sum(char.isalpha() for char in joined)
+        if SENSITIVE.search(joined) or letters < 0.6 * len(joined):
+            continue
+        paragraphs.append(joined)
+    return paragraphs
+
+
+def load_emails(limit_frac: float) -> dict[str, list[np.ndarray]]:
+    dataset = load_dataset("corbt/enron-emails", split="train").select_columns(["from", "body"])
+    if limit_frac < 1:
+        dataset = dataset.shuffle(seed=0).select(range(int(len(dataset) * limit_frac)))
+
+    parsed: list[tuple[str, list[str]]] = []
+    for row in dataset:
+        sender = (row["from"] or "").lower()
+        if "@enron.com" not in sender or BULK_SENDER.search(sender):
+            continue
+        paragraphs = clean_paragraphs(row["body"] or "")
+        if paragraphs:
+            parsed.append((sender, paragraphs))
+
+    counts = Counter(p.lower() for _, paragraphs in parsed for p in paragraphs if len(p) >= 60)
+
+    seen: set[bytes] = set()
+    splits: dict[str, list[np.ndarray]] = {"train": [], "val": [], "test": []}
+    for sender, paragraphs in parsed:
+        text = "\n\n".join(p for p in paragraphs if counts[p.lower()] < 5)
+        if not 30 <= len(text) <= 6000:
+            continue
+        digest = hashlib.sha1(re.sub(r"\s+", " ", text.lower()).encode()).digest()
+        if digest in seen:
+            continue
+        seen.add(digest)
+        bucket = int(hashlib.sha1(sender.encode()).hexdigest(), 16) % 100
+        split = "test" if bucket < 5 else "val" if bucket < 10 else "train"
+        ids = ENCODE[np.frombuffer(text.encode("ascii", "ignore"), dtype=np.uint8)]
+        splits[split].append(np.append(ids[ids != 255], EOS).astype(np.uint8))
+    return splits
+
+
+def concat_shuffled(emails: list[np.ndarray], rng: np.random.Generator) -> np.ndarray:
+    return np.concatenate([emails[i] for i in rng.permutation(len(emails))])
+
+
+def add_typos(clean: np.ndarray, word_rate: float, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
+    """Corrupt the input stream while the targets keep predicting the clean next character."""
+    x = clean[:-1].copy()
+    y = clean[1:].astype(np.int16)
+    is_letter = LETTER[x]
+    starts = np.flatnonzero(is_letter[1:] & ~is_letter[:-1]) + 1
+    ends = np.flatnonzero(is_letter[:-1] & ~is_letter[1:]) + 1
+    if len(ends) and len(starts) and ends[0] <= starts[0]:
+        ends = ends[1:]
+    count = min(len(starts), len(ends))
+    starts, ends = starts[:count], ends[:count]
+
+    lengths = ends - starts
+    chosen = (lengths >= 3) & (rng.random(count) < word_rate)
+    word_starts, word_lengths = starts[chosen], lengths[chosen]
+    position = word_starts + 1 + (rng.random(len(word_starts)) * (word_lengths - 2)).astype(np.int64)
+    position = position[position + 1 < len(x)]
+    ops = rng.integers(0, 4, size=len(position))
+    neighbor = NEIGHBOR_TABLE[x[position], rng.integers(0, 4, size=len(position))]
+
+    substitute, swap, insert, delete = (position[ops == op] for op in range(4))
+    x[substitute] = neighbor[ops == 0]
+    first, second = x[swap].copy(), x[swap + 1].copy()
+    x[swap], x[swap + 1], y[swap] = second, first, IGNORE
+
+    repeats = np.ones(len(x), dtype=np.int64)
+    repeats[delete] = 0
+    repeats[insert] = 2
+    inserted_at = np.cumsum(repeats)[insert] - 1
+    x_out, y_out = np.repeat(x, repeats), np.repeat(y, repeats)
+    x_out[inserted_at] = neighbor[ops == 2]
+    return x_out, y_out
+
+
+def make_batches(x, y, batch_size, seq_len, rng):
+    offset = int(rng.integers(seq_len))
+    starts = rng.permutation(np.arange(offset, len(x) - seq_len, seq_len))
+    span = np.arange(seq_len)
+    for index in range(0, len(starts) - batch_size + 1, batch_size):
+        positions = starts[index : index + batch_size][:, None] + span
+        yield torch.from_numpy(x[positions].astype(np.int64)), torch.from_numpy(y[positions].astype(np.int64))
+
+
+class CharGRU(nn.Module):
+    def __init__(self, hidden: int, embedding: int = 64):
+        super().__init__()
+        self.embedding = nn.Embedding(VOCAB_SIZE, embedding, padding_idx=PAD)
+        self.gru = nn.GRU(embedding, hidden, batch_first=True)
+        self.head = nn.Linear(hidden, VOCAB_SIZE)
+
+    def forward(self, ids, state=None):
+        output, state = self.gru(self.embedding(ids), state)
+        return self.head(output), state
+
+
+def sequence_loss(logits, targets, burn_in):
+    return F.cross_entropy(logits[:, burn_in:].flatten(0, 1), targets[:, burn_in:].flatten(), ignore_index=IGNORE)
+
+
+@torch.no_grad()
+def evaluate_loss(model, emails, args, device, max_chars=5_000_000):
+    rng = np.random.default_rng(1)
+    clean = concat_shuffled(emails, rng)[:max_chars]
+    x, y = clean[:-1], clean[1:].astype(np.int16)
+    total, batches = 0.0, 0
+    for inputs, targets in make_batches(x, y, args.batch, args.seq, rng):
+        logits, _ = model(inputs.to(device))
+        total += sequence_loss(logits, targets.to(device), args.burn_in).item()
+        batches += 1
+    return total / max(batches, 1)
+
+
+def decode(ids) -> str:
+    return "".join(DECODE.get(int(i), "") for i in ids)
+
+
+def sample_positions(emails, count, context, horizon, rng):
+    contexts, truths = [], []
+    attempts = 0
+    while len(contexts) < count and attempts < count * 200:
+        attempts += 1
+        email = emails[int(rng.integers(len(emails)))]
+        if len(email) < 64:
+            continue
+        cursor = int(rng.integers(16, len(email) - 16))
+        previous, following = int(email[cursor - 1]), int(email[cursor])
+        if not (LETTER[previous] or previous == ENCODE[32]) or not LETTER[following]:
+            continue
+        window = np.full(context, EOS, dtype=np.int64)
+        prefix = email[max(0, cursor - context) : cursor]
+        window[-len(prefix):] = prefix
+        truth = email[cursor : cursor + horizon]
+        ending = np.flatnonzero((truth == EOS) | (truth == NEWLINE))
+        truths.append(decode(truth[: ending[0]] if len(ending) else truth))
+        contexts.append(window)
+    return np.stack(contexts), truths
+
+
+def build_suggestion(ids, probabilities, threshold, max_words=6):
+    text, confidence, best, words = "", 1.0, "", 0
+    for token, probability in zip(ids, probabilities):
+        token = int(token)
+        if token in (EOS, NEWLINE, PAD):
+            break
+        char = DECODE[token]
+        if char == " ":
+            if confidence >= threshold:
+                best = text
+            words += 1
+            if words >= max_words:
+                break
+        confidence *= float(probability)
+        if confidence < threshold:
+            break
+        text += char
+        if char in ".,!?":
+            best = text
+    return best
+
+
+def score_suggestions(generated, probabilities, truths, thresholds):
+    report = {}
+    for threshold in thresholds:
+        shown, correct, saved = 0, 0, 0
+        for ids, probs, truth in zip(generated, probabilities, truths):
+            suggestion = build_suggestion(ids, probs, threshold)
+            if len(suggestion) < 4:
+                continue
+            shown += 1
+            if truth.startswith(suggestion):
+                correct += 1
+                saved += len(suggestion)
+        report[str(threshold)] = {
+            "coverage": round(shown / len(truths), 4),
+            "precision": round(correct / shown, 4) if shown else None,
+            "avg_chars_when_correct": round(saved / correct, 2) if correct else None,
+            "chars_saved_per_position": round(saved / len(truths), 3),
+        }
+    return report
+
+
+@torch.no_grad()
+def greedy_torch(model, contexts, steps, device, chunk=2048):
+    all_ids, all_probs = [], []
+    for index in range(0, len(contexts), chunk):
+        ids = torch.from_numpy(contexts[index : index + chunk]).to(device)
+        logits, state = model(ids)
+        step_logits = logits[:, -1]
+        batch_ids, batch_probs = [], []
+        for _ in range(steps):
+            probabilities = step_logits.softmax(-1)
+            confidence, token = probabilities.max(-1)
+            batch_ids.append(token.cpu().numpy())
+            batch_probs.append(confidence.cpu().numpy())
+            logits, state = model(token[:, None], state)
+            step_logits = logits[:, -1]
+        all_ids.append(np.stack(batch_ids, 1))
+        all_probs.append(np.stack(batch_probs, 1))
+    return np.concatenate(all_ids), np.concatenate(all_probs)
+
+
+def greedy_onnx(session, contexts, steps, hidden):
+    state = np.zeros((len(contexts), hidden), dtype=np.float32)
+    logits = None
+    for column in range(contexts.shape[1]):
+        logits, state = session.run(None, {"input_id": contexts[:, column], "state": state})
+    ids, confidences = [], []
+    for _ in range(steps):
+        probabilities = np.exp(logits - logits.max(-1, keepdims=True))
+        probabilities /= probabilities.sum(-1, keepdims=True)
+        token = probabilities.argmax(-1)
+        ids.append(token)
+        confidences.append(probabilities[np.arange(len(token)), token])
+        logits, state = session.run(None, {"input_id": token.astype(np.int64), "state": state})
+    return np.stack(ids, 1), np.stack(confidences, 1)
+
+
+class MatMulLinear(nn.Module):
+    """Linear layer that exports as MatMul so ONNX Runtime dynamic quantization applies."""
+
+    def __init__(self, weight: torch.Tensor, bias: torch.Tensor):
+        super().__init__()
+        self.weight = nn.Parameter(weight.t().clone())
+        self.bias = nn.Parameter(bias.clone())
+
+    def forward(self, inputs):
+        return inputs.matmul(self.weight) + self.bias
+
+
+class SingleStep(nn.Module):
+    """One-character GRU update with weights copied from the trained nn.GRU."""
+
+    def __init__(self, model: CharGRU, hidden: int):
+        super().__init__()
+        self.embedding = model.embedding
+        input_weight, hidden_weight = model.gru.weight_ih_l0.detach(), model.gru.weight_hh_l0.detach()
+        input_bias, hidden_bias = model.gru.bias_ih_l0.detach(), model.gru.bias_hh_l0.detach()
+        layers = []
+        for gate in range(3):
+            row = slice(gate * hidden, (gate + 1) * hidden)
+            layers.append(MatMulLinear(input_weight[row], input_bias[row]))
+            layers.append(MatMulLinear(hidden_weight[row], hidden_bias[row]))
+        self.gates = nn.ModuleList(layers)
+        self.head = MatMulLinear(model.head.weight.detach(), model.head.bias.detach())
+
+    def forward(self, input_id, state):
+        embedded = self.embedding(input_id)
+        reset = torch.sigmoid(self.gates[0](embedded) + self.gates[1](state))
+        update = torch.sigmoid(self.gates[2](embedded) + self.gates[3](state))
+        proposal = torch.tanh(self.gates[4](embedded) + reset * self.gates[5](state))
+        next_state = (1 - update) * proposal + update * state
+        return self.head(next_state), next_state
+
+
+def verify_single_step(model, step, hidden):
+    ids = torch.randint(3, VOCAB_SIZE, (4, 32))
+    logits, _ = model(ids)
+    state = torch.zeros(4, hidden)
+    outputs = []
+    for column in range(ids.shape[1]):
+        step_logits, state = step(ids[:, column], state)
+        outputs.append(step_logits)
+    difference = (torch.stack(outputs, 1) - logits).abs().max().item()
+    if difference > 1e-3:
+        raise RuntimeError(f"Exported single-step GRU does not match nn.GRU: max difference {difference}")
+    return difference
+
+
+def measure_latency(path, hidden, runs=3000):
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    session = ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"])
+    state = np.zeros((1, hidden), dtype=np.float32)
+    timings = []
+    for run in range(runs):
+        token = np.array([FIRST_PRINTABLE + run % 90], dtype=np.int64)
+        started = time.perf_counter()
+        _, state = session.run(None, {"input_id": token, "state": state})
+        timings.append((time.perf_counter() - started) * 1000)
+    return {"p50_ms": round(float(np.percentile(timings, 50)), 3), "p95_ms": round(float(np.percentile(timings, 95)), 3)}
+
+
+def export_and_score(model, args, output, contexts, truths, thresholds):
+    model = model.cpu().eval()
+    step = SingleStep(model, args.hidden).eval()
+    difference = verify_single_step(model, step, args.hidden)
+    print(f"single-step GRU matches nn.GRU (max difference {difference:.2e})", flush=True)
+
+    float_path = output / "model.onnx"
+    torch.onnx.export(
+        step,
+        (torch.zeros(1, dtype=torch.long), torch.zeros(1, args.hidden)),
+        float_path,
+        input_names=["input_id", "state"],
+        output_names=["logits", "next_state"],
+        dynamic_axes={"input_id": {0: "batch"}, "state": {0: "batch"}, "logits": {0: "batch"}, "next_state": {0: "batch"}},
+        opset_version=17,
+        dynamo=False,
+    )
+    quantized_path = output / "model_quint8.onnx"
+    quantize_dynamic(float_path, quantized_path, weight_type=QuantType.QUInt8, op_types_to_quantize=["MatMul"])
+
+    subset = min(1000, len(contexts))
+    result = {}
+    for name, path in (("float", float_path), ("quint8", quantized_path)):
+        session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"])
+        ids, probabilities = greedy_onnx(session, contexts[:subset], args.horizon, args.hidden)
+        result[name] = {
+            "size_mb": round(path.stat().st_size / 1e6, 2),
+            "native_cpu_single_thread_step": measure_latency(path, args.hidden),
+            "suggestions_1000_positions": score_suggestions(ids, probabilities, truths[:subset], thresholds),
+        }
+    return result
+
+
+def show_samples(model, device):
+    prompts = [
+        "Hi John,\n\nThanks for ", "Please let me know if ", "I will send you ", "Let me kn", "Thnaks for the ",
+        "I look forward to ", "Can we schedule a ", "Sorry for the dela", "Plese find attached ", "We need to discuss ",
+    ]
+    samples = {}
+    for prompt in prompts:
+        window = np.full(128, EOS, dtype=np.int64)
+        encoded = ENCODE[np.frombuffer(prompt.encode("ascii"), dtype=np.uint8)]
+        window[-len(encoded):] = encoded
+        ids, probabilities = greedy_torch(model, window[None], 48, device)
+        samples[prompt] = {
+            "continuation": decode(ids[0]),
+            "suggestion_at_0.2": build_suggestion(ids[0], probabilities[0], 0.2),
+        }
+        print(f"{prompt!r} -> {samples[prompt]}", flush=True)
+    return samples
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--hidden", type=int, default=512)
+    parser.add_argument("--epochs", type=int, default=30)
+    parser.add_argument("--batch", type=int, default=256)
+    parser.add_argument("--seq", type=int, default=256)
+    parser.add_argument("--burn-in", type=int, default=32)
+    parser.add_argument("--lr", type=float, default=1.5e-3)
+    parser.add_argument("--typo-rate", type=float, default=0.08)
+    parser.add_argument("--limit-frac", type=float, default=1.0)
+    parser.add_argument("--eval-positions", type=int, default=5000)
+    parser.add_argument("--context", type=int, default=128)
+    parser.add_argument("--horizon", type=int, default=48)
+    parser.add_argument("--repo", default="cfahlgren1/enron-autocomplete")
+    parser.add_argument("--no-push", action="store_true")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(0)
+    rng = np.random.default_rng(0)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"device={device} {torch.cuda.get_device_name(0) if device.type == 'cuda' else ''} args={vars(args)}", flush=True)
+    output = Path("outputs")
+    output.mkdir(parents=True, exist_ok=True)
+
+    started = time.time()
+    splits = load_emails(args.limit_frac)
+    data_stats = {name: {"emails": len(e), "chars": int(sum(map(len, e)))} for name, e in splits.items()}
+    print(f"data ready in {time.time() - started:.0f}s: {data_stats}", flush=True)
+    if not splits["train"] or not splits["val"] or not splits["test"]:
+        raise RuntimeError(f"A split is empty: {data_stats}")
+
+    model = CharGRU(args.hidden).to(device)
+    parameters = sum(p.numel() for p in model.parameters())
+    print(f"parameters={parameters:,}", flush=True)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-5)
+    steps_per_epoch = max(data_stats["train"]["chars"] // (args.batch * args.seq), 1)
+    total_steps = args.epochs * steps_per_epoch
+    scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer, lambda step: 0.02 + 0.98 * 0.5 * (1 + np.cos(np.pi * min(step / total_steps, 1)))
+    )
+
+    history = []
+    best_loss = float("inf")
+    best_path = output / "best.pt"
+    for epoch in range(args.epochs):
+        model.train()
+        epoch_started = time.time()
+        x, y = add_typos(concat_shuffled(splits["train"], rng), args.typo_rate, rng)
+        running, batches, chars = 0.0, 0, 0
+        for inputs, targets in make_batches(x, y, args.batch, args.seq, rng):
+            logits, _ = model(inputs.to(device, non_blocking=True))
+            loss = sequence_loss(logits, targets.to(device, non_blocking=True), args.burn_in)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            scheduler.step()
+            running += loss.item()
+            batches += 1
+            chars += inputs.numel()
+        if batches == 0:
+            raise RuntimeError("No training batches; reduce --batch or --seq for this data size")
+
+        model.eval()
+        val_loss = evaluate_loss(model, splits["val"], args, device)
+        seconds = time.time() - epoch_started
+        record = {
+            "epoch": epoch + 1,
+            "train_loss": round(running / batches, 4),
+            "val_loss": round(val_loss, 4),
+            "seconds": round(seconds, 1),
+            "chars_per_sec": int(chars / seconds),
+        }
+        history.append(record)
+        print(json.dumps(record), flush=True)
+        if val_loss < best_loss:
+            best_loss = val_loss
+            torch.save({"model": model.state_dict(), "args": vars(args)}, best_path)
+
+    model.load_state_dict(torch.load(best_path, map_location=device)["model"])
+    model.eval()
+
+    thresholds = [0.05, 0.1, 0.2, 0.35, 0.5, 0.7]
+    eval_rng = np.random.default_rng(42)
+    contexts, truths = sample_positions(splits["test"], args.eval_positions, args.context, args.horizon, eval_rng)
+    ids, probabilities = greedy_torch(model, contexts, args.horizon, device)
+    suggestions = score_suggestions(ids, probabilities, truths, thresholds)
+    print(json.dumps({"held_out_sender_suggestions": suggestions}, indent=1), flush=True)
+
+    examples = []
+    for index in range(min(25, len(contexts))):
+        examples.append({
+            "context": decode(contexts[index])[-60:],
+            "truth": truths[index],
+            "suggestion_at_0.2": build_suggestion(ids[index], probabilities[index], 0.2),
+        })
+    samples = show_samples(model, device)
+    exported = export_and_score(model, args, output, contexts, truths, thresholds)
+    print(json.dumps(exported, indent=1), flush=True)
+
+    report = {
+        "args": vars(args),
+        "device": str(device),
+        "parameters": parameters,
+        "data": data_stats,
+        "history": history,
+        "best_val_loss": round(best_loss, 4),
+        "test_positions": len(contexts),
+        "held_out_sender_suggestions": suggestions,
+        "exports": exported,
+        "prompt_samples": samples,
+        "held_out_examples": examples,
+        "total_seconds": round(time.time() - started, 1),
+    }
+    (output / "report.json").write_text(json.dumps(report, indent=2) + "\n")
+    (output / "vocab.json").write_text(json.dumps({
+        "pad": PAD, "eos": EOS, "newline": NEWLINE, "first_printable_id": FIRST_PRINTABLE,
+        "first_printable_codepoint": 32, "size": VOCAB_SIZE, "hidden": args.hidden,
+    }, indent=2) + "\n")
+
+    if args.no_push:
+        print(f"Skipping upload; artifacts in {output}", flush=True)
+        return
+    api = HfApi(token=os.environ["HF_TOKEN"])
+    api.create_repo(args.repo, repo_type="model", private=True, exist_ok=True)
+    api.upload_folder(folder_path=str(output), repo_id=args.repo, repo_type="model")
+    print(f"Uploaded to https://huggingface.co/{args.repo}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

vocab.json

@@ -0,0 +1,9 @@
+{
+  "pad": 0,
+  "eos": 1,
+  "newline": 2,
+  "first_printable_id": 3,
+  "first_printable_codepoint": 32,
+  "size": 98,
+  "hidden": 512
+}