Add standard Transformers and AutoTokenizer support

README.md

@@ -75,6 +75,25 @@ python inference.py \
   --prompt "Explain why the sky appears blue in two short paragraphs."
 ```
 
+For the normal Hugging Face API, load the custom architecture through the
+repository's small remote-code modules:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "North-ML1/vortex-alpha"
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo, trust_remote_code=True, torch_dtype="auto"
+)
+inputs = tokenizer("Explain why the sky appears blue.", return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=80, do_sample=False)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
+
+The `trust_remote_code=True` flag is required because Vortex's QK-Norm and
+GQA implementation is not one of the built-in Transformers model classes.
+
 For ordinary next-token completion, use the base checkpoint:
 
 ```bash

config.json

@@ -14,5 +14,10 @@
   "rms_norm_eps": 1e-5,
   "tie_word_embeddings": true,
   "attention_bias": false,
-  "torch_dtype": "bfloat16"
+  "torch_dtype": "bfloat16",
+  "use_cache": false,
+  "auto_map": {
+    "AutoConfig": "configuration_vortex.VortexConfig",
+    "AutoModelForCausalLM": "modeling_vortex.VortexForCausalLM"
+  }
 }

configuration_vortex.py

@@ -0,0 +1,47 @@
+"""Hugging Face configuration for Vortex Alpha."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class VortexConfig(PretrainedConfig):
+    model_type = "vortex"
+
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        hidden_size: int = 1024,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 4,
+        head_dim: int = 64,
+        intermediate_size: int = 3664,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 4096,
+        rope_theta: float = 100000.0,
+        rms_norm_eps: float = 1e-5,
+        attention_bias: bool = False,
+        tie_word_embeddings: bool = True,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.attention_bias = attention_bias
+        self.tie_word_embeddings = tie_word_embeddings
+        self.use_cache = use_cache
+

modeling_vortex.py

@@ -0,0 +1,255 @@
+"""Hugging Face Transformers implementation of Vortex Alpha.
+
+The model intentionally uses the same readable reference tensor names as the
+published safetensors files. It implements the full-prefix path and does not
+claim a KV cache; ``use_cache`` is therefore false by default.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from configuration_vortex import VortexConfig
+
+
+class VortexRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(hidden_states, (hidden_states.shape[-1],), self.weight, self.eps)
+
+
+class VortexRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, max_seq_len: int, theta: float) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        positions = torch.arange(max_seq_len, dtype=torch.float32)
+        frequencies = torch.outer(positions, inv_freq)
+        angles = torch.cat((frequencies, frequencies), dim=-1)
+        self.register_buffer("cos_cached", angles.cos()[None, None], persistent=False)
+        self.register_buffer("sin_cached", angles.sin()[None, None], persistent=False)
+
+    @staticmethod
+    def rotate_half(x: torch.Tensor) -> torch.Tensor:
+        half = x.shape[-1] // 2
+        return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if position_ids is None:
+            cos = self.cos_cached[:, :, : query.shape[-2]]
+            sin = self.sin_cached[:, :, : query.shape[-2]]
+        else:
+            cos = self.cos_cached[0, 0, position_ids].unsqueeze(1)
+            sin = self.sin_cached[0, 0, position_ids].unsqueeze(1)
+        cos = cos.to(dtype=query.dtype, device=query.device)
+        sin = sin.to(dtype=query.dtype, device=query.device)
+        return (
+            query * cos + self.rotate_half(query) * sin,
+            key * cos + self.rotate_half(key) * sin,
+        )
+
+
+class VortexAttention(nn.Module):
+    def __init__(self, config: VortexConfig) -> None:
+        super().__init__()
+        if config.num_attention_heads % config.num_key_value_heads:
+            raise ValueError("num_attention_heads must divide evenly by num_key_value_heads")
+        if config.num_attention_heads * config.head_dim != config.hidden_size:
+            raise ValueError("num_attention_heads * head_dim must equal hidden_size")
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        kv_dim = config.num_key_value_heads * config.head_dim
+        self.q_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+        self.k_proj = nn.Linear(config.hidden_size, kv_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(config.hidden_size, kv_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+        self.q_norm = VortexRMSNorm(config.head_dim, config.rms_norm_eps)
+        self.k_norm = VortexRMSNorm(config.head_dim, config.rms_norm_eps)
+        self.rope = VortexRotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch, seq_len, _ = hidden_states.shape
+        query = self.q_proj(hidden_states).view(batch, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key = self.k_proj(hidden_states).view(batch, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value = self.v_proj(hidden_states).view(batch, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        query, key = self.rope(self.q_norm(query), self.k_norm(key), position_ids)
+        if self.num_key_value_heads != self.num_heads:
+            repeats = self.num_heads // self.num_key_value_heads
+            key = key.repeat_interleave(repeats, dim=1)
+            value = value.repeat_interleave(repeats, dim=1)
+
+        sdpa_mask = None
+        is_causal = attention_mask is None and seq_len > 1
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                valid_keys = attention_mask.to(dtype=torch.bool, device=hidden_states.device)[:, None, None, :]
+                causal = torch.ones((seq_len, seq_len), dtype=torch.bool, device=hidden_states.device).tril()
+                sdpa_mask = valid_keys & causal[None, None, :, :]
+            elif attention_mask.ndim == 4:
+                sdpa_mask = attention_mask
+            else:
+                raise ValueError("Vortex expects a 2-D or 4-D attention mask")
+        output = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=sdpa_mask,
+            dropout_p=0.0,
+            is_causal=is_causal,
+        )
+        output = output.transpose(1, 2).contiguous().view(batch, seq_len, -1)
+        return self.o_proj(output)
+
+
+class VortexMLP(nn.Module):
+    def __init__(self, config: VortexConfig) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.attention_bias)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.attention_bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))
+
+
+class VortexDecoderBlock(nn.Module):
+    def __init__(self, config: VortexConfig) -> None:
+        super().__init__()
+        self.norm1 = VortexRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attn = VortexAttention(config)
+        self.norm2 = VortexRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.ffn = VortexMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask, position_ids)
+        hidden_states = hidden_states + self.ffn(self.norm2(hidden_states))
+        return hidden_states
+
+
+class VortexPreTrainedModel(PreTrainedModel):
+    config_class = VortexConfig
+    base_model_prefix = "vortex"
+    supports_gradient_checkpointing = False
+
+
+class VortexForCausalLM(VortexPreTrainedModel):
+    _tied_weights_keys = ["embed_tokens.weight"]
+    main_input_name = "input_ids"
+
+    def __init__(self, config: VortexConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(VortexDecoderBlock(config) for _ in range(config.num_hidden_layers))
+        self.norm = VortexRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self._initialize_weights()
+
+    def _initialize_weights(self) -> None:
+        output_std = 0.02 / math.sqrt(2.0 * self.config.num_hidden_layers)
+        nn.init.normal_(self.embed_tokens.weight, mean=0.0, std=0.02)
+        for block in self.layers:
+            for module in block.modules():
+                if isinstance(module, VortexRMSNorm):
+                    nn.init.ones_(module.weight)
+                elif isinstance(module, nn.Linear):
+                    is_output = module is block.attn.o_proj or module is block.ffn.down_proj
+                    nn.init.normal_(module.weight, mean=0.0, std=output_std if is_output else 0.02)
+                    if module.bias is not None:
+                        nn.init.zeros_(module.bias)
+        nn.init.ones_(self.norm.weight)
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.embed_tokens
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.embed_tokens = new_embeddings
+
+    def prepare_inputs_for_generation(self, input_ids: torch.Tensor, **kwargs) -> dict:
+        # No KV cache is exported. Returning the full prefix keeps generation
+        # correct, though slower than a cache-enabled implementation.
+        return {
+            "input_ids": input_ids,
+            "attention_mask": kwargs.get("attention_mask"),
+            "use_cache": False,
+        }
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("pass either input_ids or inputs_embeds, not both")
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise ValueError("input_ids or inputs_embeds is required")
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            hidden_states = inputs_embeds
+        if position_ids is None:
+            position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
+        all_hidden_states = () if output_hidden_states else None
+        for block in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            hidden_states = block(hidden_states, attention_mask, position_ids)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+        hidden_states = self.norm(hidden_states)
+        logits = F.linear(hidden_states, self.embed_tokens.weight)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous().float()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+        if return_dict is False:
+            output = (logits, None, all_hidden_states, None)
+            return ((loss,) + output) if loss is not None else output
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=None,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+

tokenization_vortex.py

@@ -0,0 +1,78 @@
+"""SentencePiece tokenizer wrapper for Hugging Face AutoTokenizer."""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+
+from transformers import PreTrainedTokenizer
+
+
+class VortexTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.model"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        unk_token: str = "<unk>",
+        eos_token: str = "</s>",
+        **kwargs,
+    ) -> None:
+        import sentencepiece as spm
+
+        self.vocab_file = str(vocab_file)
+        self.sp_model = spm.SentencePieceProcessor(model_file=self.vocab_file)
+        super().__init__(unk_token=unk_token, eos_token=eos_token, **kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.sp_model.vocab_size())
+
+    def get_vocab(self) -> dict[str, int]:
+        vocab = {self.sp_model.id_to_piece(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text: str, **kwargs) -> list[str]:
+        return list(self.sp_model.encode(text, out_type=str))
+
+    def _convert_token_to_id(self, token: str) -> int:
+        token_id = int(self.sp_model.piece_to_id(token))
+        return self.unk_token_id if token_id < 0 else token_id
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.sp_model.id_to_piece(int(index))
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return self.sp_model.decode_pieces(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is None:
+            return list(token_ids_0)
+        return list(token_ids_0) + list(token_ids_1)
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        if already_has_special_tokens:
+            return [0] * len(token_ids_0)
+        total = len(token_ids_0) + (len(token_ids_1) if token_ids_1 is not None else 0)
+        return [0] * total
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+        directory = Path(save_directory)
+        directory.mkdir(parents=True, exist_ok=True)
+        name = (filename_prefix + "-" if filename_prefix else "") + "tokenizer.model"
+        output = directory / name
+        if Path(self.vocab_file).resolve() != output.resolve():
+            shutil.copyfile(self.vocab_file, output)
+        return (str(output),)
+

tokenizer_config.json

@@ -1,6 +1,9 @@
 {
   "model_type": "vortex",
-  "tokenizer_class": "SentencePieceProcessor",
+  "tokenizer_class": "VortexTokenizer",
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_vortex.VortexTokenizer", null]
+  },
   "vocab_size": 8192,
   "unk_token": "<unk>",
   "eos_token": "</s>",