Upload modular/code/modular_data.py with huggingface_hub

modular/code/modular_data.py

@@ -0,0 +1,112 @@
+"""
+Modular arithmetic dataset — a + b mod p.
+
+Matches Nanda et al. (2023) exactly: p=113, 30% train split, fixed seed.
+
+Token layout (contiguous integer IDs, no external tokenizer):
+  0 … p-1  →  numbers
+  p        →  '+'
+  p+1      →  '='
+  p+2      →  PAD
+  total trajectory vocab = p+3
+
+Sequence format: [a, +, b, =, result]  (5 tokens, prompt_len=4)
+"""
+import json
+import random
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import List, Tuple
+
+P: int = 113  # Nanda's prime
+
+PLUS   = P        # 113
+EQUALS = P + 1    # 114
+PAD    = P + 2    # 115
+VOCAB_SIZE = P + 3  # 116  — trajectory vocab passed to SorlModelWrapper
+
+PROMPT_LEN = 4   # [a, +, b, =]
+ANSWER_LEN = 1   # [result]
+SEQ_LEN    = 5
+
+EVAL_CACHE_DIR = Path(__file__).resolve().parent / "eval_sets"
+
+
+@dataclass
+class ModularExample:
+    tokens: List[int]   # [a, +, b, =, result]
+    a: int
+    b: int
+    result: int
+
+
+def make_example(a: int, b: int, p: int = P) -> ModularExample:
+    return ModularExample(tokens=[a, PLUS, b, EQUALS, (a + b) % p], a=a, b=b, result=(a + b) % p)
+
+
+def generate_dataset(p: int = P, train_fraction: float = 0.3, seed: int = 42) -> Tuple[List[ModularExample], List[ModularExample]]:
+    """All p² pairs shuffled and split into fixed train/test."""
+    rng = random.Random(seed)
+    all_pairs = [(a, b) for a in range(p) for b in range(p)]
+    rng.shuffle(all_pairs)
+    n_train = int(len(all_pairs) * train_fraction)
+    train = [make_example(a, b, p) for a, b in all_pairs[:n_train]]
+    test  = [make_example(a, b, p) for a, b in all_pairs[n_train:]]
+    return train, test
+
+
+DATASET_REPO = "thoughtworks/arithmetic-sorl-data"
+
+
+def get_eval_set(p: int = P, seed: int = 42) -> List[ModularExample]:
+    """Load test set — local cache → HF download → regenerate."""
+    cache = EVAL_CACHE_DIR / f"modular_p{p}_test_seed{seed}.json"
+    if cache.exists():
+        with open(cache) as f:
+            return [ModularExample(**ex) for ex in json.load(f)]
+
+    EVAL_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    try:
+        from huggingface_hub import hf_hub_download
+        path = hf_hub_download(
+            repo_id=DATASET_REPO, repo_type="dataset",
+            filename=f"modular/test_seed{seed}.json",
+        )
+        import shutil
+        shutil.copy(path, cache)
+        with open(cache) as f:
+            return [ModularExample(**ex) for ex in json.load(f)]
+    except Exception:
+        pass
+
+    _, test = generate_dataset(p=p, seed=seed)
+    with open(cache, "w") as f:
+        json.dump([asdict(ex) for ex in test], f)
+    return test
+
+
+def get_train_set(p: int = P, seed: int = 42) -> List[ModularExample]:
+    """Load train set — local cache → HF download → regenerate."""
+    cache = EVAL_CACHE_DIR / f"modular_p{p}_train_seed{seed}.json"
+    if cache.exists():
+        with open(cache) as f:
+            return [ModularExample(**ex) for ex in json.load(f)]
+
+    EVAL_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    try:
+        from huggingface_hub import hf_hub_download
+        path = hf_hub_download(
+            repo_id=DATASET_REPO, repo_type="dataset",
+            filename=f"modular/train_seed{seed}.json",
+        )
+        import shutil
+        shutil.copy(path, cache)
+        with open(cache) as f:
+            return [ModularExample(**ex) for ex in json.load(f)]
+    except Exception:
+        pass
+
+    train, _ = generate_dataset(p=p, seed=seed)
+    with open(cache, "w") as f:
+        json.dump([asdict(ex) for ex in train], f)
+    return train