Upload modular/code/evaluate.py with huggingface_hub

modular/code/evaluate.py

@@ -0,0 +1,62 @@
+"""
+Modular arithmetic evaluator — single-token accuracy via recursion.
+
+Usage:
+    from arithmetic.modular.training.evaluate import ModularEvaluator
+    ev = ModularEvaluator(model, device="cuda", K=1)
+    acc = ev.run(test_examples)
+"""
+import torch
+from typing import List
+
+from sorl.sorl_trainer import infer_insert_mask, insert_tokens_with_padding, expand_prompt_len
+from arithmetic.modular.data.modular import ModularExample, PROMPT_LEN, PAD
+
+
+class ModularEvaluator:
+
+    def __init__(self, model, device: str = "cuda", K: int = 1):
+        self.model  = model
+        self.device = device
+        self.K      = K
+        self.base_v = int(model.vocab_sizes[0].item())
+
+    @torch.no_grad()
+    def run(self, examples: List[ModularExample], max_examples: int = 0) -> float:
+        self.model.eval()
+        if max_examples > 0:
+            examples = examples[:max_examples]
+        correct = 0
+        for ex in examples:
+            ids  = torch.tensor(ex.tokens, dtype=torch.long, device=self.device).unsqueeze(0)
+            attn = torch.ones_like(ids)
+            pl   = torch.tensor([PROMPT_LEN], dtype=torch.long, device=self.device)
+
+            im = infer_insert_mask(ids, self.K, attn)
+            ep = expand_prompt_len(pl, im)
+            ed, ea = insert_tokens_with_padding(ids, attn, im, self.base_v, PAD)
+
+            data, _, _ = self.model.recursion(
+                ed, ea, max_iterations=2,
+                memory_span_abs=512, memory_span_traj=512,
+                temperature=0.0, prompt_len=ep,
+            )
+
+            # Forward pass to get logits on the filled sequence
+            block_mask = self.model._create_sorl_block_mask(data, 512, 512)
+            out = self.model.model.forward(
+                input_ids=data, attention_mask=ea,
+                block_mask=block_mask, use_cache=False,
+            )
+            logits = out.logits
+
+            # Result token is the (PROMPT_LEN)-th trajectory token (0-indexed)
+            is_traj   = data[0] < self.base_v
+            traj_pos  = is_traj.nonzero(as_tuple=True)[0]
+            result_pos = traj_pos[PROMPT_LEN].item()
+            pred = logits[0, result_pos - 1, :self.base_v].argmax().item()
+
+            if pred == ex.result:
+                correct += 1
+
+        return correct / max(len(examples), 1)