Fix tokenizer and sampling edge cases and make asset downloads atomic

.github/workflows/ci.yml

@@ -21,6 +21,8 @@ jobs:
         with:
           components: rustfmt, clippy
       - uses: Swatinem/rust-cache@v2
+      - name: asset downloader tests (no network)
+        run: python3 -m unittest discover -s python -p 'test_*.py'
       - name: rustfmt
         run: cargo fmt --all -- --check
       - name: clippy

docs/correctness.md

@@ -93,3 +93,30 @@ exist:
 - No SSM `scan` op (not implemented).
 - No full transformer-LM end-to-end parity (no LM yet) — only the MLP is verified
   end-to-end, plus every individual operator above.
+
+## Invalid-input contracts and sampling
+
+CPU reference operators check dimension products before allocation. Normalization
+requires a positive width and a finite non-negative epsilon; rotary embeddings
+require a finite positive theta. Attention validates head widths and bounds
+absolute causal offsets without wrapping.
+
+Use Sampler::sample_token for checked single-token sampling, and Sampler::validate
+to check temperature. Empty or non-finite logits are rejected. Probabilities use
+f64 intermediates; top-k ties prefer smaller token IDs and the retained candidates
+have a stable ID order. Seeded sampled sequences may differ from earlier versions
+because their candidate ordering and arithmetic were corrected. Greedy finite
+argmax retains its earliest-ID tie behavior.
+
+Tokenizer merge parsing retains hash-prefixed merge tokens, rejects malformed or
+duplicate pairs, and checks both operands. Repeated whitespace preserves GPT-2's
+lookahead boundary before the next token. The additive try_decode and
+try_decode_bytes APIs report unknown token IDs; existing decode APIs retain their
+permissive behavior.
+
+The download script accepts --out, --timeout, and --force. Each transfer uses a
+temporary file beside its destination, checks any advertised Content-Length, and
+atomically replaces the destination only after completion. Existing nonempty
+assets are skipped by default; --force refreshes them. Network-free downloader
+regressions run in CI. These tests do not download or validate pretrained GPT-2
+weights and do not replace on-device Metal parity tests.

python/fetch_gpt2.py

@@ -1,29 +1,69 @@
-"""Download GPT-2 (124M) weights + tokenizer into models/gpt2/.
-
-Pulls the public HuggingFace `openai-community/gpt2` files (~550 MB) with no
-extra dependencies beyond the standard library. These files are gitignored.
-"""
-
+"""Download GPT-2 weights and tokenizer with atomic, bounded transfers."""
+import argparse
+import math
 import os
+import shutil
+import sys
+import tempfile
 import urllib.request
 
 BASE = "https://huggingface.co/openai-community/gpt2/resolve/main"
 FILES = ["model.safetensors", "vocab.json", "merges.txt", "config.json"]
 OUT = os.path.join("models", "gpt2")
 
 
-def main():
-    os.makedirs(OUT, exist_ok=True)
-    for name in FILES:
-        dst = os.path.join(OUT, name)
-        if os.path.exists(dst) and os.path.getsize(dst) > 0:
-            print(f"  have {name}")
-            continue
-        print(f"  downloading {name} …")
-        urllib.request.urlretrieve(f"{BASE}/{name}", dst)
-    print(f"Done. Weights in {OUT}/")
-    print('Try:  cargo run --release --bin generate -- --prompt "Once upon a time"')
+def positive_timeout(value):
+    try:
+        seconds = float(value)
+    except ValueError as error:
+        raise argparse.ArgumentTypeError("timeout must be a number") from error
+    if not math.isfinite(seconds) or seconds <= 0:
+        raise argparse.ArgumentTypeError("timeout must be finite and positive")
+    return seconds
+
+
+def download_file(url, destination, timeout=60.0):
+    """Install a completed transfer atomically, preserving any previous asset."""
+    directory = os.path.dirname(os.path.abspath(destination))
+    descriptor, temporary = tempfile.mkstemp(prefix=".gpt2-", suffix=".part", dir=directory)
+    os.close(descriptor)
+    try:
+        with urllib.request.urlopen(url, timeout=timeout) as response, open(temporary, "wb") as output:
+            shutil.copyfileobj(response, output)
+            expected = response.headers.get("Content-Length")
+        size = os.path.getsize(temporary)
+        if size == 0:
+            raise OSError("downloaded file is empty")
+        if expected is not None and size != int(expected):
+            raise OSError(f"incomplete transfer: expected {expected} bytes, received {size}")
+        os.replace(temporary, destination)
+    finally:
+        if os.path.exists(temporary):
+            os.unlink(temporary)
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--out", default=OUT, help="asset directory (default: models/gpt2)")
+    parser.add_argument("--timeout", type=positive_timeout, default=60.0, help="socket timeout in seconds")
+    parser.add_argument("--force", action="store_true", help="replace existing assets")
+    args = parser.parse_args(argv)
+    try:
+        os.makedirs(args.out, exist_ok=True)
+        for name in FILES:
+            destination = os.path.join(args.out, name)
+            if not args.force and os.path.isfile(destination) and os.path.getsize(destination) > 0:
+                print(f"  have {name}")
+                continue
+            print(f"  downloading {name} …")
+            download_file(f"{BASE}/{name}", destination, args.timeout)
+    except (OSError, ValueError) as error:
+        print(f"Download failed: {error}", file=sys.stderr)
+        return 1
+    print(f"Done. Weights in {args.out}/")
+    print(f'Try: cargo run --release --bin generate -- --model-dir "{args.out}" --prompt "Once upon a time"')
+    return 0
 
 
 if __name__ == "__main__":
-    main()
+    raise SystemExit(main())

python/test_fetch_gpt2.py

@@ -0,0 +1,74 @@
+"""Network-free regression tests for GPT-2 asset installation."""
+import argparse
+import io
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import fetch_gpt2
+
+
+class Response(io.BytesIO):
+    def __init__(self, payload=b"complete", length=None):
+        super().__init__(payload)
+        self.headers = {} if length is None else {"Content-Length": str(length)}
+
+
+class DownloadTests(unittest.TestCase):
+    def test_complete_download_replaces_asset(self):
+        with tempfile.TemporaryDirectory() as directory:
+            destination = os.path.join(directory, "model")
+            with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(length=8)) as opener:
+                fetch_gpt2.download_file("https://example.test/model", destination, 12.0)
+                opener.assert_called_once_with("https://example.test/model", timeout=12.0)
+            with open(destination, "rb") as stream:
+                self.assertEqual(stream.read(), b"complete")
+            self.assertEqual(os.listdir(directory), ["model"])
+
+    def test_failure_preserves_previous_asset(self):
+        with tempfile.TemporaryDirectory() as directory:
+            destination = os.path.join(directory, "model")
+            with open(destination, "wb") as stream:
+                stream.write(b"previous")
+            with patch("fetch_gpt2.urllib.request.urlopen", side_effect=OSError("interrupted")):
+                with self.assertRaises(OSError):
+                    fetch_gpt2.download_file("https://example.test/model", destination)
+            with open(destination, "rb") as stream:
+                self.assertEqual(stream.read(), b"previous")
+            self.assertEqual(os.listdir(directory), ["model"])
+
+    def test_empty_and_truncated_downloads_never_become_final_assets(self):
+        for payload, length in [(b"", None), (b"short", 100)]:
+            with tempfile.TemporaryDirectory() as directory:
+                destination = os.path.join(directory, "model")
+                with patch("fetch_gpt2.urllib.request.urlopen", return_value=Response(payload, length)):
+                    with self.assertRaises(OSError):
+                        fetch_gpt2.download_file("https://example.test/model", destination)
+                self.assertEqual(os.listdir(directory), [])
+
+    def test_timeout_validation(self):
+        for value in ["0", "-1", "NaN", "inf", "invalid"]:
+            with self.assertRaises(argparse.ArgumentTypeError):
+                fetch_gpt2.positive_timeout(value)
+        self.assertEqual(fetch_gpt2.positive_timeout("5"), 5.0)
+
+    def test_existing_assets_are_skipped_unless_forced(self):
+        with tempfile.TemporaryDirectory() as directory:
+            for name in fetch_gpt2.FILES:
+                with open(os.path.join(directory, name), "wb") as stream:
+                    stream.write(b"previous")
+            with patch("builtins.print"), patch("fetch_gpt2.download_file") as transfer:
+                self.assertEqual(fetch_gpt2.main(["--out", directory]), 0)
+                transfer.assert_not_called()
+                self.assertEqual(fetch_gpt2.main(["--out", directory, "--force", "--timeout", "3"]), 0)
+                self.assertEqual(transfer.call_count, len(fetch_gpt2.FILES))
+
+    def test_main_reports_network_failure_as_unsuccessful(self):
+        with tempfile.TemporaryDirectory() as directory:
+            with patch("builtins.print"), patch("fetch_gpt2.download_file", side_effect=OSError("network")):
+                self.assertEqual(fetch_gpt2.main(["--out", directory]), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()