Add standard Transformers inference adapter

README.md

@@ -35,7 +35,25 @@ What is Python?<|im_end|>
 <|im_start|>assistant
 ```
 
-The model is a custom Aurora checkpoint, not a Transformers-compatible architecture. Use the included native Aurora runtime, YAML config, and tokenizer.
+The model is a custom Aurora checkpoint. The included native Aurora runtime is the simplest path; a Transformers remote-code adapter is also provided below for normal Hub-style testing.
+
+## Transformers / Hugging Face test
+
+The repository also includes a Transformers remote-code adapter, so it can be loaded through the normal `AutoTokenizer` and `AutoModelForCausalLM` APIs:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "North-ML1/Aurora-Proelia-ChatML"
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
+messages = [{"role": "user", "content": "What is Python?"}]
+inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+outputs = model.generate(inputs, max_new_tokens=96, do_sample=False)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+`trust_remote_code=True` is required because Aurora is a custom architecture; inspect the repository code before enabling it in an untrusted environment.
 
 ## What changed
 

config.json

@@ -1,13 +1,33 @@
 {
   "_name_or_path": "North-ML1/Aurora-Proelia-ChatML",
-  "architectures": [
-    "AuroraForCausalLM"
-  ],
+  "architectures": ["AuroraForCausalLM"],
+  "auto_map": {
+    "AutoConfig": "configuration_aurora.AuroraHFConfig",
+    "AutoModelForCausalLM": "modeling_aurora.AuroraForCausalLM"
+  },
   "model_type": "aurora",
   "model_name": "Aurora Proelia ChatML",
   "vocab_size": 16000,
+  "hidden_size": 896,
+  "num_layers": 23,
+  "num_attention_heads": 14,
+  "num_key_value_heads": 2,
+  "intermediate_size": 2432,
+  "context_length": 2048,
   "max_position_embeddings": 2048,
+  "rope_theta": 500000.0,
+  "rms_norm_eps": 0.00001,
+  "qk_norm": true,
+  "tie_word_embeddings": true,
+  "attention_bias": false,
+  "mlp_bias": false,
+  "dropout": 0.0,
+  "num_experts": 1,
+  "router_aux_loss_coef": 0.0,
+  "router_z_loss_coef": 0.0,
+  "router_noise_scale": 0.0,
+  "moe_capacity_factor": 0.0,
+  "router_use_gate_weight": false,
   "torch_dtype": "float16",
-  "library_name": "aurora",
-  "chat_template": "{% for message in messages %}<|im_start|>{{ message['role'] }}\\n{{ message['content'] }}<|im_end|>\\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\\n{% endif %}"
-}\n
+  "library_name": "aurora"
+}

configuration_aurora.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class AuroraHFConfig(PretrainedConfig):
+    model_type = "aurora"
+
+    def __init__(
+        self,
+        model_name: str = "Aurora Proelia ChatML",
+        vocab_size: int = 16000,
+        hidden_size: int = 896,
+        num_layers: int = 23,
+        num_attention_heads: int = 14,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 2432,
+        context_length: int = 2048,
+        rope_theta: float = 500000.0,
+        rms_norm_eps: float = 1.0e-5,
+        qk_norm: bool = True,
+        tie_word_embeddings: bool = True,
+        attention_bias: bool = False,
+        mlp_bias: bool = False,
+        dropout: float = 0.0,
+        num_experts: int = 1,
+        router_aux_loss_coef: float = 0.0,
+        router_z_loss_coef: float = 0.0,
+        router_noise_scale: float = 0.0,
+        moe_capacity_factor: float = 0.0,
+        router_use_gate_weight: bool = False,
+        **kwargs,
+    ):
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            bos_token_id=kwargs.pop("bos_token_id", 1),
+            eos_token_id=kwargs.pop("eos_token_id", 2),
+            pad_token_id=kwargs.pop("pad_token_id", 0),
+            **kwargs,
+        )
+        self.model_name = model_name
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.intermediate_size = intermediate_size
+        self.context_length = context_length
+        self.max_position_embeddings = context_length
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.qk_norm = qk_norm
+        self.tie_word_embeddings = tie_word_embeddings
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.dropout = dropout
+        self.num_experts = num_experts
+        self.router_aux_loss_coef = router_aux_loss_coef
+        self.router_z_loss_coef = router_z_loss_coef
+        self.router_noise_scale = router_noise_scale
+        self.moe_capacity_factor = moe_capacity_factor
+        self.router_use_gate_weight = router_use_gate_weight
+        self.use_cache = False

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "do_sample": false,
+  "max_new_tokens": 96,
+  "use_cache": false
+}

modeling_aurora.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_aurora import AuroraHFConfig
+from .aurora.config import AuroraConfig as NativeAuroraConfig
+from .aurora.model import AuroraForCausalLM as NativeAuroraForCausalLM
+
+
+class AuroraForCausalLM(PreTrainedModel):
+    config_class = AuroraHFConfig
+    base_model_prefix = "aurora"
+    main_input_name = "input_ids"
+    _supports_cache_class = False
+
+    def __init__(self, config: AuroraHFConfig):
+        super().__init__(config)
+        native_config = NativeAuroraConfig(
+            model_name=config.model_name,
+            vocab_size=config.vocab_size,
+            hidden_size=config.hidden_size,
+            num_layers=config.num_layers,
+            num_attention_heads=config.num_attention_heads,
+            num_key_value_heads=config.num_key_value_heads,
+            intermediate_size=config.intermediate_size,
+            context_length=config.context_length,
+            rope_theta=config.rope_theta,
+            rms_norm_eps=config.rms_norm_eps,
+            qk_norm=config.qk_norm,
+            tie_word_embeddings=config.tie_word_embeddings,
+            attention_bias=config.attention_bias,
+            mlp_bias=config.mlp_bias,
+            dropout=config.dropout,
+            num_experts=config.num_experts,
+            router_aux_loss_coef=config.router_aux_loss_coef,
+            router_z_loss_coef=config.router_z_loss_coef,
+            router_noise_scale=config.router_noise_scale,
+            moe_capacity_factor=config.moe_capacity_factor,
+            router_use_gate_weight=config.router_use_gate_weight,
+        )
+        self.aurora = NativeAuroraForCausalLM(native_config)
+
+    def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False, **kwargs):
+        # The raw Aurora export stores the native module at the repository root;
+        # the Transformers wrapper stores it below `aurora`.
+        remapped = {
+            (key if key.startswith("aurora.") else f"aurora.{key}"): value
+            for key, value in state_dict.items()
+        }
+        return super().load_state_dict(remapped, strict=strict, assign=assign, **kwargs)
+
+    def get_input_embeddings(self):
+        return self.aurora.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.aurora.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.aurora.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.aurora.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        logits, loss = self.aurora(input_ids=input_ids, labels=labels)
+        return CausalLMOutputWithPast(loss=loss, logits=logits)

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "pad_token": "<pad>"
+}

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "bos_token": "<bos>",
+  "eos_token": "<eos>",
+  "pad_token": "<pad>",
+  "model_max_length": 2048,
+  "clean_up_tokenization_spaces": false,
+  "chat_template": "{% for message in messages %}<|im_start|>{{ message['role'] }}\\n{{ message['content'] }}<|im_end|>\\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\\n{% endif %}"
+}