Use NeroXSAForCausalLM architecture name

README.md

@@ -106,7 +106,7 @@ text = generate(
 print(text)
 ```
 
-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`.
+Nero XS 2 uses the custom PyTorch class `NeroXSAForCausalLM`. 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 | `NeroXSA11ForCausalLM` |
+| Architecture class | `NeroXSAForCausalLM` |
 | Stored parameters | **2,993,152** |
 | Physical / effective blocks | 10 / 14 |
 | Recurrent layout | 1 prelude + 4 middle blocks × 2 passes + 5 coda |

config.json

@@ -1,6 +1,6 @@
 {
-  "architectures": ["NeroXSA11ForCausalLM"],
-  "model_type": "nero_xsa11",
+  "architectures": ["NeroXSAForCausalLM"],
+  "model_type": "nero_xsa",
   "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 NeroXSA11ForCausalLM(nn.Module):
+class NeroXSAForCausalLM(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 NeroXSA11ForCausalLM(variant_config(name))
+    return NeroXSAForCausalLM(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 = NeroXSA11ForCausalLM(cfg)
+    model = NeroXSAForCausalLM(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: NeroXSA11ForCausalLM,
+    model: NeroXSAForCausalLM,
     tokenizer,
     prompt: str,
     max_new_tokens: int = 64,
@@ -338,8 +338,9 @@ def generate(
     return tokenizer.decode(ids, skip_special_tokens=True)
 
 
-# Backward-compatible import alias for the initial release name.
-NeroXSA2ForCausalLM = NeroXSA11ForCausalLM
+# Backward-compatible import aliases for the earlier release names.
+NeroXSA11ForCausalLM = NeroXSAForCausalLM
+NeroXSA2ForCausalLM = NeroXSAForCausalLM
 
 
 if __name__ == "__main__":