Fix arithmetic generation cache default

README.md

@@ -80,9 +80,34 @@ text = "12 + 34 ="
 inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
 
 with torch.no_grad():
-    outputs = model(**inputs)
+    output_ids = model.generate(
+        **inputs,
+        max_new_tokens=3,
+        do_sample=False,
+    )
+
+print(tokenizer.decode(output_ids[0], skip_special_tokens=True))
+# 12 + 34 = 46
 ```
 
+### Generation Cache
+
+Atom2.7m derives arithmetic `place_ids` and `role_ids` from the full token sequence. During arithmetic generation, result digits need those features to be recomputed from the full current prefix. For this reason, the checkpoint defaults to `use_cache=False` for `model.generate(...)`. This is slower than KV-cache generation, but preserves the arithmetic feature annotations used by the model.
+
+You can opt into faster cached generation when exact arithmetic-aware generation is not required:
+
+```python
+output_ids = model.generate(
+    **inputs,
+    max_new_tokens=32,
+    use_cache=True,
+)
+```
+
+Cached generation remains supported, but for arithmetic continuations it may
+produce lower-quality results unless the caller supplies correctly updated
+`place_ids` and `role_ids`.
+
 ## Evaluation
 
 ### ArithMark 2.0
@@ -121,6 +146,12 @@ implementation can score this slightly longer harness window; reduce batch size
 or set `max_length` to the longest sequence found if a task variant contains
 longer continuations.
 
+For multiple-choice or benchmark-style evaluation, no special generation cache
+setting is required. Log-likelihood scoring runs full `context + continuation`
+forward passes, so arithmetic features are derived from the complete sequence.
+This is the path used by the included ArithMark benchmark script and by
+lm-evaluation-harness log-likelihood tasks.
+
 ## Results
 
 | Benchmark | Metric | Value |

config.json

@@ -68,6 +68,7 @@
   "role_vocab_size": 12,
   "rope_theta": 5000.0,
   "transformers_version": "4.57.6",
+  "use_cache": false,
   "use_place_embeddings": true,
   "use_role_embeddings": true,
   "vocab_size": 4096,

generation_config.json

@@ -1,4 +1,5 @@
 {
   "_from_model_config": true,
-  "transformers_version": "4.57.6"
+  "transformers_version": "4.57.6",
+  "use_cache": false
 }

model.py

@@ -341,7 +341,8 @@ class GPTForCausalLM(GPTPreTrainedModel, GenerationMixin):
         return embeddings
 
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
-        if past_key_values is not None and past_key_values.get_seq_length() > 0:
+        use_cache = bool(kwargs.get("use_cache", False))
+        if use_cache and past_key_values is not None and past_key_values.get_seq_length() > 0:
             input_ids = input_ids[:, -1:]
             if kwargs.get("place_ids") is not None:
                 kwargs["place_ids"] = kwargs["place_ids"][:, -1:]
@@ -352,8 +353,8 @@ class GPTForCausalLM(GPTPreTrainedModel, GenerationMixin):
             "place_ids": kwargs.get("place_ids"),
             "role_ids": kwargs.get("role_ids"),
             "attention_mask": attention_mask,
-            "past_key_values": past_key_values,
-            "use_cache": True,
+            "past_key_values": past_key_values if use_cache else None,
+            "use_cache": use_cache,
         }
 
     def _get_freqs_cis(self, seq_len, device):