Upload folder using huggingface_hub

README.md

@@ -15,3 +15,18 @@ library_name: pytorch
 From-scratch narrow-domain coding SLM for React + Tailwind generation and unified-diff edits.
 
 Includes trained `model.safetensors` weights.
+
+## Transformers Usage
+
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+model_id = "Sharjeelbaig/neurocoder"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+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)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```

config.json

@@ -1,17 +1,29 @@
 {
   "architectures": [
-    "TinyMoEModel"
+    "NeuroCoderForCausalLM"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_neurocoder.NeuroCoderConfig",
+    "AutoModelForCausalLM": "modeling_neurocoder.NeuroCoderForCausalLM",
+    "AutoTokenizer": [
+      "tokenization_neurocoder.NeuroCoderTokenizer",
+      null
+    ]
+  },
+  "bos_token_id": 1,
   "capacity_factor_infer": 1.0,
   "capacity_factor_train": 1.25,
   "context_length": 320,
+  "eos_token_id": 2,
   "ffn_multiplier": 4,
   "hidden_size": 256,
-  "model_type": "tinymoe",
+  "model_type": "neurocoder",
   "moe_every_n_layers": 2,
   "num_experts": 4,
   "num_heads": 8,
   "num_layers": 8,
+  "pad_token_id": 0,
   "top_k": 2,
+  "unk_token_id": 3,
   "vocab_size": 1714
 }

configuration_neurocoder.py

@@ -0,0 +1,50 @@
+"""Transformers config for NeuroCoder remote-code loading."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class NeuroCoderConfig(PretrainedConfig):
+    model_type = "neurocoder"
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        context_length: int = 4096,
+        hidden_size: int = 1024,
+        num_layers: int = 20,
+        num_heads: int = 16,
+        ffn_multiplier: int = 4,
+        moe_every_n_layers: int = 2,
+        num_experts: int = 8,
+        top_k: int = 2,
+        capacity_factor_train: float = 1.25,
+        capacity_factor_infer: float = 1.0,
+        dropout: float = 0.0,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.context_length = context_length
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        # Aliases expected by Transformers generation/runtime utilities.
+        self.num_hidden_layers = num_layers
+        self.num_attention_heads = num_heads
+        self.max_position_embeddings = context_length
+        self.use_cache = False
+        self.ffn_multiplier = ffn_multiplier
+        self.moe_every_n_layers = moe_every_n_layers
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.capacity_factor_train = capacity_factor_train
+        self.capacity_factor_infer = capacity_factor_infer
+        self.dropout = dropout
+
+    @property
+    def head_dim(self) -> int:
+        if self.hidden_size % self.num_heads != 0:
+            raise ValueError("hidden_size must be divisible by num_heads")
+        return self.hidden_size // self.num_heads

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ec3e8e9e20e7cc41eaacb419be91939f3b8ff25e494afe5ba063f2efdf50a1e6
+oid sha256:662bfd3a3fabe2977d92c697faaa0af70c6704d5581fd9549d578a994e13202a
 size 75081480

modeling_neurocoder.py

@@ -0,0 +1,224 @@
+"""Transformers model implementation for NeuroCoder remote-code loading."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_neurocoder import NeuroCoderConfig
+except Exception:
+    from configuration_neurocoder import NeuroCoderConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, x: Tensor) -> Tensor:
+        rms = x.pow(2).mean(-1, keepdim=True)
+        return x * torch.rsqrt(rms + self.eps) * self.weight
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, config: NeuroCoderConfig) -> None:
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.head_dim = config.head_dim
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3)
+        self.out = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, x: Tensor) -> Tensor:
+        bsz, seq_len, hidden = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        def shape_heads(t: Tensor) -> Tensor:
+            return t.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        q = shape_heads(q)
+        k = shape_heads(k)
+        v = shape_heads(v)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+        mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))
+        attn = attn.masked_fill(~mask, float("-inf"))
+        probs = F.softmax(attn, dim=-1)
+        out = torch.matmul(probs, v)
+        out = out.transpose(1, 2).contiguous().view(bsz, seq_len, hidden)
+        return self.out(out)
+
+
+class DenseFFN(nn.Module):
+    def __init__(self, config: NeuroCoderConfig) -> None:
+        super().__init__()
+        inner = config.hidden_size * config.ffn_multiplier
+        self.gate = nn.Linear(config.hidden_size, inner)
+        self.up = nn.Linear(config.hidden_size, inner)
+        self.down = nn.Linear(inner, config.hidden_size)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+
+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.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)
+        self.experts = nn.ModuleList([DenseFFN(config) for _ in range(config.num_experts)])
+
+    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        bsz, seq_len, hidden = x.shape
+        x_flat = x.reshape(-1, hidden)
+        tokens = x_flat.shape[0]
+
+        logits = self.router(x_flat)
+        probs = F.softmax(logits, dim=-1)
+        top_vals, top_idx = torch.topk(probs, k=self.top_k, dim=-1)
+
+        capacity_factor = self.capacity_factor_train if self.training else self.capacity_factor_infer
+        capacity = max(1, math.ceil(capacity_factor * tokens / self.num_experts))
+
+        output = torch.zeros_like(x_flat)
+        expert_load = []
+
+        for expert_id in range(self.num_experts):
+            expert = self.experts[expert_id]
+            assigned_indices = []
+            assigned_weights = []
+            for rank in range(self.top_k):
+                mask = top_idx[:, rank] == expert_id
+                idx = torch.nonzero(mask, as_tuple=False).squeeze(-1)
+                if idx.numel() == 0:
+                    continue
+                weights = top_vals[idx, rank]
+                assigned_indices.append(idx)
+                assigned_weights.append(weights)
+
+            if not assigned_indices:
+                expert_load.append(0.0)
+                continue
+
+            token_indices = torch.cat(assigned_indices, dim=0)
+            token_weights = torch.cat(assigned_weights, dim=0)
+            if token_indices.numel() > capacity:
+                token_indices = token_indices[:capacity]
+                token_weights = token_weights[:capacity]
+
+            expert_in = x_flat[token_indices]
+            expert_out = expert(expert_in)
+            output[token_indices] += expert_out * token_weights.unsqueeze(-1)
+            expert_load.append(float(token_indices.numel() / max(tokens, 1)))
+
+        load_tensor = torch.tensor(expert_load, device=x.device)
+        mean_prob = probs.mean(dim=0)
+        aux_loss = self.num_experts * torch.sum(mean_prob * load_tensor)
+        z_loss = torch.mean(torch.logsumexp(logits, dim=-1) ** 2)
+        return output.reshape(bsz, seq_len, hidden), aux_loss, z_loss
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: NeuroCoderConfig, use_moe: bool) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(config.hidden_size)
+        self.norm2 = RMSNorm(config.hidden_size)
+        self.attn = SelfAttention(config)
+        self.ffn = MoEFeedForward(config) if use_moe else DenseFFN(config)
+        self.use_moe = use_moe
+
+    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+        x = x + self.attn(self.norm1(x))
+        aux_loss = torch.tensor(0.0, device=x.device)
+        z_loss = torch.tensor(0.0, device=x.device)
+        ffn_input = self.norm2(x)
+        if self.use_moe:
+            ffn_out, aux_loss, z_loss = self.ffn(ffn_input)
+        else:
+            ffn_out = self.ffn(ffn_input)
+        x = x + ffn_out
+        return x, aux_loss, z_loss
+
+
+class NeuroCoderForCausalLM(PreTrainedModel):
+    config_class = NeuroCoderConfig
+    base_model_prefix = "neurocoder"
+    _no_split_modules = ["TransformerBlock", "MoEFeedForward"]
+
+    def __init__(self, config: NeuroCoderConfig) -> None:
+        super().__init__(config)
+        self.token_embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(config, use_moe=((idx + 1) % config.moe_every_n_layers == 0))
+                for idx in range(config.num_layers)
+            ]
+        )
+        self.norm = RMSNorm(config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.token_embed.weight
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.token_embed
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.token_embed = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Tensor,
+        **kwargs: Any,
+    ) -> dict[str, Tensor]:
+        return {"input_ids": input_ids}
+
+    def forward(
+        self,
+        input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        labels: Tensor | None = None,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast:
+        if input_ids is None:
+            raise ValueError("input_ids is required")
+
+        x = self.token_embed(input_ids)
+        aux_loss = torch.tensor(0.0, device=input_ids.device)
+        z_loss = torch.tensor(0.0, device=input_ids.device)
+
+        for layer in self.layers:
+            x, layer_aux, layer_z = layer(x)
+            aux_loss = aux_loss + layer_aux
+            z_loss = z_loss + layer_z
+
+        x = self.norm(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+                ignore_index=-100,
+            )
+            loss = loss + 0.01 * aux_loss + 0.001 * z_loss
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits)

special_tokens_map.json

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

tokenization_neurocoder.py

@@ -0,0 +1,79 @@
+"""Transformers tokenizer for NeuroCoder remote-code loading."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import re
+from typing import Any
+
+from transformers import PreTrainedTokenizer
+
+TOKEN_PATTERN = re.compile(r"\s+|[A-Za-z_][A-Za-z0-9_]*|\d+|\S")
+SPECIAL_TOKENS = ["<pad>", "<bos>", "<eos>", "<unk>"]
+
+
+class NeuroCoderTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(self, vocab_file: str | None = None, **kwargs: Any) -> None:
+        self.vocab: dict[str, int] = {}
+        self.id_to_token: list[str] = []
+
+        if vocab_file is not None:
+            payload = json.loads(Path(vocab_file).read_text(encoding="utf-8"))
+            self.vocab = {str(k): int(v) for k, v in payload.get("vocab", {}).items()}
+            max_id = max(self.vocab.values()) if self.vocab else -1
+            self.id_to_token = ["<unk>"] * (max_id + 1)
+            for token, idx in self.vocab.items():
+                self.id_to_token[idx] = token
+
+        if not self.vocab:
+            self.vocab = {token: idx for idx, token in enumerate(SPECIAL_TOKENS)}
+            self.id_to_token = SPECIAL_TOKENS[:]
+
+        kwargs.setdefault("bos_token", "<bos>")
+        kwargs.setdefault("eos_token", "<eos>")
+        kwargs.setdefault("unk_token", "<unk>")
+        kwargs.setdefault("pad_token", "<pad>")
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    def get_vocab(self) -> dict[str, int]:
+        return dict(self.vocab)
+
+    def _tokenize(self, text: str) -> list[str]:
+        return TOKEN_PATTERN.findall(text)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.vocab.get(token, self.vocab.get(self.unk_token, 0))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        if 0 <= index < len(self.id_to_token):
+            return self.id_to_token[index]
+        return self.unk_token
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return "".join(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 token_ids_0
+        return token_ids_0 + token_ids_1
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+        out_dir = Path(save_directory)
+        out_dir.mkdir(parents=True, exist_ok=True)
+        file_name = "tokenizer.json" if filename_prefix is None else f"{filename_prefix}-tokenizer.json"
+        out_path = out_dir / file_name
+        payload = {
+            "type": "simple_regex_tokenizer",
+            "special_tokens": SPECIAL_TOKENS,
+            "vocab": self.vocab,
+        }
+        out_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+        return (str(out_path),)

tokenizer.json

@@ -1,4 +1,5 @@
 {
+  "added_tokens": [],
   "special_tokens": [
     "<pad>",
     "<bos>",

tokenizer_config.json

@@ -1,4 +1,10 @@
 {
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_neurocoder.NeuroCoderTokenizer",
+      null
+    ]
+  },
   "model_max_length": 320,
   "padding_side": "right",
   "special_tokens_map": {