Update app.py

app.py

@@ -1,6 +1,84 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PretrainedConfig
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+# Model architecture definition
+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 SmolLM2ForCausalLM(PreTrainedModel):
+    config_class = SmolLM2Config
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        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):
+        hidden_states = self.embed_tokens(input_ids)
+        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))
+            
+        return logits if loss is None else (loss, logits)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+# Register the model architecture
+from transformers import AutoConfig, AutoModelForCausalLM
+AutoConfig.register("smollm2", SmolLM2Config)
+AutoModelForCausalLM.register(SmolLM2Config, SmolLM2ForCausalLM)
 
 # Load model and tokenizer
 model_id = "jatingocodeo/SmolLM2"