Update app.py

app.py

@@ -1,231 +1,8 @@
 import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig
-import torch.nn as nn
-import torch.nn.functional as F
-import math
-
-class SmolLM2Config(PretrainedConfig):
-    model_type = "smollm2"
-    
-    def __init__(
-        self,
-        vocab_size=49152,
-        hidden_size=576,
-        intermediate_size=1536,
-        num_hidden_layers=30,
-        num_attention_heads=9,
-        num_key_value_heads=3,
-        hidden_act="silu",
-        max_position_embeddings=2048,
-        initializer_range=0.041666666666666664,
-        rms_norm_eps=1e-5,
-        use_cache=True,
-        pad_token_id=None,
-        bos_token_id=0,
-        eos_token_id=0,
-        tie_word_embeddings=True,
-        rope_theta=10000.0,
-        **kwargs
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_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.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        super().__init__(
-            pad_token_id=pad_token_id,
-            bos_token_id=bos_token_id,
-            eos_token_id=eos_token_id,
-            tie_word_embeddings=tie_word_embeddings,
-            **kwargs
-        )
-
-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()
-        
-        # Project and reshape
-        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)
-        
-        # Repeat k/v heads if needed
-        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)
-        
-        # Transpose for attention
-        q = q.transpose(1, 2)  # (batch, num_heads, seq_len, head_dim)
-        k = k.transpose(1, 2)  # (batch, num_heads, seq_len, head_dim)
-        v = v.transpose(1, 2)  # (batch, num_heads, seq_len, head_dim)
-        
-        # Calculate attention scores
-        scale = 1.0 / math.sqrt(self.head_dim)
-        scores = torch.matmul(q, k.transpose(-2, -1)) * scale  # (batch, num_heads, seq_len, seq_len)
-        
-        # Apply attention mask if provided
-        if attention_mask is not None:
-            # Ensure mask is broadcastable
-            if attention_mask.dim() == 2:
-                attention_mask = attention_mask.unsqueeze(1).unsqueeze(1)  # (batch, 1, 1, seq_len)
-            scores = scores + attention_mask
-        
-        # Apply softmax and dropout
-        attention_weights = F.softmax(scores, dim=-1)
-        
-        # Apply attention to values
-        output = torch.matmul(attention_weights, v)  # (batch, num_heads, seq_len, head_dim)
-        
-        # Reshape and project back
-        output = output.transpose(1, 2).contiguous()  # (batch, seq_len, num_heads, head_dim)
-        output = output.view(batch_size, seq_length, -1)  # (batch, seq_len, hidden_size)
-        output = self.o_proj(output)
-        
-        return output
-
-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 SmolLM2ForCausalLM(PreTrainedModel):
-    config_class = SmolLM2Config
-    _no_split_modules = ["LlamaDecoderLayer"]
-    
-    def __init__(self, config):
-        super().__init__(config)
-        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:
-            self.lm_head.weight = self.embed_tokens.weight
-            
-    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=None, **kwargs):
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        
-        hidden_states = self.embed_tokens(input_ids)
-        
-        # Create causal attention mask
-        batch_size, seq_length = input_ids.size()
-        device = input_ids.device
-        
-        # Create causal mask
-        causal_mask = torch.triu(
-            torch.ones((seq_length, seq_length), dtype=torch.bool, device=device),
-            diagonal=1
-        )
-        causal_mask = causal_mask.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, seq_len]
-        causal_mask = causal_mask.expand(batch_size, 1, seq_length, seq_length)
-        causal_mask = causal_mask.to(dtype=torch.float32) * -1e4
-        
-        # Combine with attention mask if provided
-        if attention_mask is not None:
-            # Convert attention mask to float and unsqueeze
-            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)  # [batch, 1, 1, seq_len]
-            attention_mask = attention_mask.expand(batch_size, 1, seq_length, seq_length)
-            attention_mask = (1.0 - attention_mask) * -1e4
-            # Combine masks
-            causal_mask = causal_mask + attention_mask
-        
-        # Process through layers
-        for layer in self.layers:
-            hidden_states = layer(hidden_states, causal_mask)
-            
-        hidden_states = self.norm(hidden_states)
-        logits = self.lm_head(hidden_states)
-        
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))
-        
-        if return_dict:
-            from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
-            return CausalLMOutputWithCrossAttentions(
-                loss=loss,
-                logits=logits,
-                past_key_values=None,
-                hidden_states=None,
-                attentions=None,
-                cross_attentions=None,
-            )
-        return (loss, logits) if loss is not None else logits
-
-    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
-        # Only return what we need
-        inputs = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask
-        }
-        return inputs
-
-# Register the model architecture
-from transformers import AutoConfig
-AutoConfig.register("smollm2", SmolLM2Config)
-AutoModelForCausalLM.register(SmolLM2Config, SmolLM2ForCausalLM)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download
+import json
 
 # Cache for model and tokenizer
 MODEL = None
@@ -239,10 +16,12 @@ def initialize():
         model_id = "jatingocodeo/SmolLM2"
         
         try:
+            # Download model files from HF Hub
+            config_path = hf_hub_download(repo_id=model_id, filename="config.json")
+            
             # Load tokenizer
-            print("\n1. Loading tokenizer...")
+            print("Loading tokenizer...")
             TOKENIZER = AutoTokenizer.from_pretrained(model_id)
-            print("✓ Tokenizer loaded successfully")
             
             # Add special tokens if needed
             special_tokens = {
@@ -250,25 +29,25 @@ def initialize():
                 'eos_token': '</s>',
                 'bos_token': '<s>'
             }
-            num_added = TOKENIZER.add_special_tokens(special_tokens)
-            print(f"✓ Added {num_added} special tokens")
+            TOKENIZER.add_special_tokens(special_tokens)
             
             # Load model
-            print("\n2. Loading model...")
+            print("Loading model...")
             MODEL = AutoModelForCausalLM.from_pretrained(
                 model_id,
-                trust_remote_code=True,
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                trust_remote_code=True,
                 low_cpu_mem_usage=True
             )
             
-            # Move model to appropriate device
+            # Move model to device
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            MODEL = MODEL.to(device)
-            print(f"✓ Model loaded successfully and moved to {device}")
+            MODEL.to(device)
+            
+            print(f"Model loaded successfully on {device}")
             
         except Exception as e:
-            print(f"Error initializing model: {str(e)}")
+            print(f"Error initializing: {str(e)}")
             raise
 
 def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
@@ -281,6 +60,7 @@ def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
         if not prompt.strip():
             return "Please enter a prompt."
         
+        # Add BOS token if needed
         if not prompt.startswith(TOKENIZER.bos_token):
             prompt = TOKENIZER.bos_token + prompt
         
@@ -288,36 +68,26 @@ def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
         input_ids = TOKENIZER.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
         input_ids = input_ids.to(MODEL.device)
         
-        # Create attention mask
-        attention_mask = torch.ones_like(input_ids)
-        
         # Generate
         with torch.no_grad():
-            output_ids = MODEL.generate(
+            outputs = MODEL.generate(
                 input_ids,
-                attention_mask=attention_mask,
                 max_length=min(max_length + len(input_ids[0]), 2048),
                 temperature=temperature,
                 top_k=top_k,
                 do_sample=True,
                 pad_token_id=TOKENIZER.pad_token_id,
                 eos_token_id=TOKENIZER.eos_token_id,
-                num_return_sequences=1,
-                use_cache=True
+                num_return_sequences=1
             )
         
         # Decode and return
-        generated_text = TOKENIZER.decode(output_ids[0], skip_special_tokens=True)
+        generated_text = TOKENIZER.decode(outputs[0], skip_special_tokens=True)
         return generated_text.strip()
         
     except Exception as e:
-        print(f"Generation error details: {str(e)}")
-        import traceback
-        traceback.print_exc()
-        return f"Error generating text: {str(e)}"
-
-# Initialize on startup
-initialize()
+        print(f"Error generating text: {str(e)}")
+        return f"An error occurred: {str(e)}"
 
 # Create Gradio interface
 iface = gr.Interface(