release: hfstd v2 stabilization and inference suite

README.md

@@ -27,6 +27,6 @@ model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
 
 prompt = "Generate a landing page for marketing agency titled Velocity Landing"
 inputs = tokenizer(prompt, return_tensors="pt")
-outputs = model.generate(**inputs, max_new_tokens=220, temperature=0.7, do_sample=True)
+outputs = model.generate(**inputs, max_new_tokens=220, temperature=0.7, do_sample=True, use_cache=False)
 print(tokenizer.decode(outputs[0], skip_special_tokens=True))
 ```

config.json

@@ -23,7 +23,7 @@
   "num_heads": 8,
   "num_layers": 8,
   "pad_token_id": 0,
-  "top_k": 2,
+  "router_top_k": 2,
   "unk_token_id": 3,
   "vocab_size": 1714
 }

configuration_neurocoder.py

@@ -18,6 +18,7 @@ class NeuroCoderConfig(PretrainedConfig):
         ffn_multiplier: int = 4,
         moe_every_n_layers: int = 2,
         num_experts: int = 8,
+        router_top_k: int | None = None,
         top_k: int = 2,
         capacity_factor_train: float = 1.25,
         capacity_factor_infer: float = 1.0,
@@ -38,7 +39,8 @@ class NeuroCoderConfig(PretrainedConfig):
         self.ffn_multiplier = ffn_multiplier
         self.moe_every_n_layers = moe_every_n_layers
         self.num_experts = num_experts
-        self.top_k = top_k
+        # Keep MoE router top-k separate from generation top_k to avoid HF generation warnings.
+        self.router_top_k = router_top_k if router_top_k is not None else top_k
         self.capacity_factor_train = capacity_factor_train
         self.capacity_factor_infer = capacity_factor_infer
         self.dropout = dropout

modeling_neurocoder.py

@@ -74,7 +74,7 @@ class MoEFeedForward(nn.Module):
     def __init__(self, config: NeuroCoderConfig) -> None:
         super().__init__()
         self.num_experts = config.num_experts
-        self.top_k = config.top_k
+        self.top_k = config.router_top_k
         self.capacity_factor_train = config.capacity_factor_train
         self.capacity_factor_infer = config.capacity_factor_infer
         self.router = nn.Linear(config.hidden_size, config.num_experts, bias=False)