Rename architecture class to NeroXSA11ForCausalLM

README.md

@@ -17,7 +17,7 @@ datasets:
 - HuggingFaceFW/finephrase
 - HuggingFaceTB/finemath
 model-index:
-- name: Nero XS 1.1
+- name: Nero XS 2
   results:
   - task:
       type: text-generation
@@ -106,7 +106,7 @@ text = generate(
 print(text)
 ```
 
-Nero XS 2 uses the custom PyTorch class `NeroXSA2ForCausalLM`. It is not a drop-in Transformers `AutoModelForCausalLM` checkpoint. The complete standalone implementation is included in `modeling_nero_xs_2.py`.
+Nero XS 2 uses the custom PyTorch class `NeroXSA11ForCausalLM`. It is not a drop-in Transformers `AutoModelForCausalLM` checkpoint. The complete standalone implementation is included in `modeling_nero_xs_2.py`.
 
 ## Results
 
@@ -129,7 +129,7 @@ HellaSwag, ARC, and PIQA were evaluated with `lm-eval` 0.4.12 using exact contin
 
 | Component | Configuration |
 |---|---|
-| Architecture class | `NeroXSA2ForCausalLM` |
+| Architecture class | `NeroXSA11ForCausalLM` |
 | Stored parameters | **2,993,152** |
 | Physical / effective blocks | 10 / 14 |
 | Recurrent layout | 1 prelude + 4 middle blocks × 2 passes + 5 coda |
@@ -240,4 +240,4 @@ The released safetensors file was produced directly from the portable final chec
 
 ## License
 
-The original Nero XS 1.1 weights and repository material are released under CC-BY-4.0. Attribution is required. Upstream dataset terms and attribution requirements remain applicable to their respective source material.
+The original Nero XS 2 weights and repository material are released under CC-BY-4.0. Attribution is required. Upstream dataset terms and attribution requirements remain applicable to their respective source material.

config.json

@@ -1,6 +1,6 @@
 {
-  "architectures": ["NeroXSA2ForCausalLM"],
-  "model_type": "nero_xs_2",
+  "architectures": ["NeroXSA11ForCausalLM"],
+  "model_type": "nero_xsa11",
   "vocab_size": 2048,
   "width": 128,
   "heads": 4,

modeling_nero_xs_2.py

@@ -199,7 +199,7 @@ class ReleasedNeroXSControl(nn.Module):
         return F.linear(self.norm(x), self.embedding.weight)
 
 
-class NeroXSA2ForCausalLM(nn.Module):
+class NeroXSA11ForCausalLM(nn.Module):
     """XSA + selective recurrence + loop conditioning + EngramLite."""
 
     def __init__(self, cfg: NeroConfig = NeroConfig()):
@@ -257,7 +257,7 @@ def variant_config(name: str) -> NeroConfig:
 def build_variant(name: str) -> nn.Module:
     if name == "released_control":
         return ReleasedNeroXSControl()
-    return NeroXSA2ForCausalLM(variant_config(name))
+    return NeroXSA11ForCausalLM(variant_config(name))
 
 
 def architecture_audit() -> dict[str, int | float | bool]:
@@ -290,7 +290,7 @@ def load_model(model_dir: str | Path, device: str | torch.device = "cpu"):
     raw = json.loads((model_dir / "config.json").read_text())
     fields = NeroConfig.__dataclass_fields__
     cfg = NeroConfig(**{key: value for key, value in raw.items() if key in fields})
-    model = NeroXSA2ForCausalLM(cfg)
+    model = NeroXSA11ForCausalLM(cfg)
     model.load_state_dict(load_file(model_dir / "model.safetensors"), strict=True)
     model.to(device).eval()
     tokenizer = AutoTokenizer.from_pretrained(model_dir)
@@ -299,7 +299,7 @@ def load_model(model_dir: str | Path, device: str | torch.device = "cpu"):
 
 @torch.inference_mode()
 def generate(
-    model: NeroXSA2ForCausalLM,
+    model: NeroXSA11ForCausalLM,
     tokenizer,
     prompt: str,
     max_new_tokens: int = 64,
@@ -338,6 +338,10 @@ def generate(
     return tokenizer.decode(ids, skip_special_tokens=True)
 
 
+# Backward-compatible import alias for the initial release name.
+NeroXSA2ForCausalLM = NeroXSA11ForCausalLM
+
+
 if __name__ == "__main__":
     import json
     print("NERO_XS2_AUDIT=" + json.dumps(architecture_audit(), sort_keys=True))