Upload src/model.py with huggingface_hub

src/model.py

@@ -0,0 +1,199 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional
+from config import NexusConfig
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.float() * rms * self.weight.float()).type_as(x)
+
+def precompute_freqs_cis(config: NexusConfig) -> torch.Tensor:
+    dim = config.dim // config.num_heads
+    freqs = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2).float() / dim))
+    t = torch.arange(config.max_seq_len)
+    freqs = torch.outer(t, freqs)
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, config: NexusConfig):
+        super().__init__()
+        self.freqs_cis = precompute_freqs_cis(config)
+
+    def forward(self, x: torch.Tensor, start_pos: int = 0):
+        _, seq_len, _, head_dim = x.shape
+        freqs_cis = self.freqs_cis[start_pos:start_pos+seq_len, :head_dim//2].to(x.device)
+        freqs_cis = freqs_cis.view(1, seq_len, 1, head_dim//2)
+        
+        x_shaped = x.float().reshape(*x.shape[:-1], -1, 2)
+        x_complex = torch.complex(x_shaped[..., 0], x_shaped[..., 1])
+        
+        x_rotated = x_complex * freqs_cis
+        
+        x_out = torch.stack([x_rotated.real, x_rotated.imag], dim=-1).reshape_as(x_shaped)
+        
+        return x_out.reshape_as(x).type_as(x)
+
+class Attention(nn.Module):
+    def __init__(self, config: NexusConfig):
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.num_kv_heads = config.num_kv_heads
+        if self.num_kv_heads is None:
+            self.num_kv_heads = config.num_heads
+        self.head_dim = config.dim // config.num_heads
+        self.num_kv_groups = config.num_heads // self.num_kv_heads
+
+        self.wq = nn.Linear(config.dim, config.dim, bias=False)
+        self.wk = nn.Linear(config.dim, self.head_dim * self.num_kv_heads, bias=False)
+        self.wv = nn.Linear(config.dim, self.head_dim * self.num_kv_heads, bias=False)
+        self.wo = nn.Linear(config.dim, config.dim, bias=False)
+        self.rotary = RotaryEmbedding(config)
+
+    def forward(self, x: torch.Tensor, start_pos: int = 0, mask: Optional[torch.Tensor] = None):
+        bsz, seqlen, _ = x.shape
+
+        q = self.wq(x).view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q = self.rotary(q, start_pos)
+        k = self.rotary(k, start_pos)
+
+        if self.num_kv_groups > 1:
+            k = k[:, :, None, :, :].expand(bsz, self.num_kv_heads, self.num_kv_groups, seqlen, self.head_dim)
+            k = k.reshape(bsz, self.num_heads, seqlen, self.head_dim)
+            v = v[:, :, None, :, :].expand(bsz, self.num_kv_heads, self.num_kv_groups, seqlen, self.head_dim)
+            v = v.reshape(bsz, self.num_heads, seqlen, self.head_dim)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale
+
+        if mask is not None:
+            attn_weights = attn_weights + mask
+
+        attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(q)
+        attn_output = torch.matmul(attn_weights, v)
+
+        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        return self.wo(attn_output)
+
+class FeedForward(nn.Module):
+    def __init__(self, config: NexusConfig):
+        super().__init__()
+        hidden_dim = int(2 * config.ff_dim / 3)
+        hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
+
+        self.w1 = nn.Linear(config.dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, config.dim, bias=False)
+        self.w3 = nn.Linear(config.dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: NexusConfig):
+        super().__init__()
+        self.attention = Attention(config)
+        self.feed_forward = FeedForward(config)
+        self.attention_norm = RMSNorm(config.dim, config.norm_eps)
+        self.ff_norm = RMSNorm(config.dim, config.norm_eps)
+
+    def forward(self, x: torch.Tensor, start_pos: int = 0, mask: Optional[torch.Tensor] = None):
+        h = x + self.attention(self.attention_norm(x), start_pos, mask)
+        out = h + self.feed_forward(self.ff_norm(h))
+        return out
+
+class Nexus(nn.Module):
+    def __init__(self, config: NexusConfig):
+        super().__init__()
+        self.config = config
+
+        self.token_embeddings = nn.Embedding(config.vocab_size, config.dim)
+        self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_layers)])
+        self.norm = RMSNorm(config.dim, config.norm_eps)
+        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+
+        self.token_embeddings.weight = self.output.weight
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids: torch.Tensor, start_pos: int = 0):
+        _, seqlen = input_ids.shape
+
+        mask = torch.full((1, 1, seqlen, start_pos + seqlen), float('-inf'),
+                          dtype=torch.float32, device=input_ids.device)
+        mask = torch.triu(mask, diagonal=start_pos + 1).type_as(input_ids)
+
+        x = self.token_embeddings(input_ids)
+
+        for layer in self.layers:
+            x = layer(x, start_pos, mask)
+
+        x = self.norm(x)
+        logits = self.output(x)
+
+        return logits
+
+    def generate(self, input_ids: torch.Tensor, max_new_tokens: int,
+                 temperature: float = 0.7, top_k: int = 50, top_p: float = 0.9):
+        self.eval()
+        generated = []
+
+        for _ in range(max_new_tokens):
+            seq_len = input_ids.shape[1]
+            if seq_len > self.config.max_seq_len:
+                input_ids = input_ids[:, -self.config.max_seq_len:]
+
+            with torch.no_grad():
+                logits = self(input_ids, 0)
+                logits = logits[:, -1, :] / temperature
+
+                if top_k > 0:
+                    top_k_values, _ = torch.topk(logits, top_k)
+                    min_top_k = top_k_values[:, -1].unsqueeze(-1)
+                    logits = torch.where(logits < min_top_k,
+                                         torch.full_like(logits, float('-inf')), logits)
+
+                if top_p > 0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[:, 0] = False
+                    indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
+                    indices_to_remove = indices_to_remove.scatter(1, sorted_indices,
+                                                                  sorted_indices_to_remove)
+                    logits = torch.where(indices_to_remove,
+                                         torch.full_like(logits, float('-inf')), logits)
+
+                probs = F.softmax(logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            generated.append(next_token.item())
+
+        return generated, input_ids
+
+def create_nexus_model():
+    from config import nexus_config
+    config = nexus_config
+
+    model = Nexus(config)
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print(f"[Nexus SmAll] Model created with {total_params/1e6:.1f}M parameters "
+          f"({trainable_params/1e6:.1f}M trainable)")
+
+    return model, config