Update app.py

app.py

@@ -5,7 +5,92 @@ import torch.nn as nn
 import torch.nn.functional as F
 import math
 
-# Model architecture definition
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return self.weight * x
+
+class LlamaAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+
+    def forward(self, hidden_states, attention_mask=None):
+        batch_size, seq_length, _ = hidden_states.size()
+        
+        q = self.q_proj(hidden_states).view(batch_size, seq_length, self.num_heads, self.head_dim)
+        k = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
+        v = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
+        
+        if self.num_kv_heads < self.num_heads:
+            k = k.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
+            v = v.repeat_interleave(self.num_heads // self.num_kv_heads, dim=2)
+        
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        
+        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+            
+        attention_probs = F.softmax(attention_scores, dim=-1)
+        context = torch.matmul(attention_probs, v)
+        
+        context = context.transpose(1, 2).contiguous()
+        context = context.view(batch_size, seq_length, -1)
+        
+        return self.o_proj(context)
+
+class LlamaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        gate = self.act_fn(self.gate_proj(x))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+class LlamaDecoderLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self_attn = LlamaAttention(config)
+        self.mlp = LlamaMLP(config)
+        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+    def forward(self, hidden_states, attention_mask=None):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.self_attn(hidden_states, attention_mask)
+        hidden_states = residual + hidden_states
+        
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        return hidden_states
+
 class SmolLM2Config(PretrainedConfig):
     model_type = "smollm2"
     
@@ -58,6 +143,8 @@ class SmolLM2ForCausalLM(PreTrainedModel):
         self.config = config
         
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         
         if config.tie_word_embeddings:
@@ -65,6 +152,20 @@ class SmolLM2ForCausalLM(PreTrainedModel):
             
     def forward(self, input_ids, attention_mask=None, labels=None):
         hidden_states = self.embed_tokens(input_ids)
+        
+        # Create causal attention mask if none provided
+        if attention_mask is None:
+            attention_mask = torch.triu(
+                torch.ones((input_ids.size(1), input_ids.size(1)), dtype=torch.bool, device=input_ids.device),
+                diagonal=1
+            )
+            attention_mask = attention_mask.unsqueeze(0).unsqueeze(0)
+            attention_mask = attention_mask * -1e4
+        
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask)
+            
+        hidden_states = self.norm(hidden_states)
         logits = self.lm_head(hidden_states)
         
         loss = None
@@ -74,7 +175,10 @@ class SmolLM2ForCausalLM(PreTrainedModel):
         return logits if loss is None else (loss, logits)
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {"input_ids": input_ids}
+        return {
+            "input_ids": input_ids,
+            "attention_mask": kwargs.get("attention_mask", None)
+        }
 
 # Register the model architecture
 from transformers import AutoConfig, AutoModelForCausalLM