Update app.py

app.py

@@ -4,6 +4,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
+import os
 
 class RMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-5):
@@ -190,37 +191,56 @@ model_id = "jatingocodeo/SmolLM2"
 
 def load_model():
     try:
+        print("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(model_id)
+        print("Tokenizer loaded successfully")
+        
         # Ensure the tokenizer has the necessary special tokens
         special_tokens = {
             'pad_token': '[PAD]',
             'eos_token': '</s>',
             'bos_token': '<s>'
         }
+        print("Adding special tokens...")
         tokenizer.add_special_tokens(special_tokens)
         
-        # Load model without device_map
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            pad_token_id=tokenizer.pad_token_id
+        print("Loading model configuration...")
+        config = SmolLM2Config()
+        
+        print("Initializing model...")
+        model = SmolLM2ForCausalLM(config)
+        
+        print("Loading model weights...")
+        state_dict = torch.load(
+            os.path.join(model_id, "pytorch_model.bin"),
+            map_location="cpu"
         )
+        model.load_state_dict(state_dict)
         
         # Move model to device manually
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Moving model to device: {device}")
         model = model.to(device)
         
         # Resize token embeddings to match new tokenizer
+        print("Resizing token embeddings...")
         model.resize_token_embeddings(len(tokenizer))
+        
+        print("Model loaded successfully!")
         return model, tokenizer
     except Exception as e:
         print(f"Error loading model: {str(e)}")
+        print(f"Error type: {type(e)}")
+        import traceback
+        traceback.print_exc()
         raise
 
 def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
     try:
+        print(f"\nGenerating text for prompt: {prompt}")
         # Load model and tokenizer (caching them for subsequent calls)
         if not hasattr(generate_text, "model"):
+            print("First call - loading model...")
             generate_text.model, generate_text.tokenizer = load_model()
         
         # Ensure the prompt is not empty
@@ -231,15 +251,17 @@ def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
         if not prompt.startswith(generate_text.tokenizer.bos_token):
             prompt = generate_text.tokenizer.bos_token + prompt
         
+        print("Encoding prompt...")
         # Encode the prompt
         input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=2048)
         input_ids = input_ids.to(generate_text.model.device)
         
+        print("Generating text...")
         # Generate text
         with torch.no_grad():
             output_ids = generate_text.model.generate(
                 input_ids,
-                max_length=min(max_length + len(input_ids[0]), 2048),  # Respect model's max length
+                max_length=min(max_length + len(input_ids[0]), 2048),
                 temperature=temperature,
                 top_k=top_k,
                 do_sample=True,
@@ -248,12 +270,17 @@ def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
                 num_return_sequences=1
             )
         
+        print("Decoding generated text...")
         # Decode and return the generated text
         generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        print("Generation completed successfully!")
         return generated_text.strip()
     
     except Exception as e:
         print(f"Error during generation: {str(e)}")
+        print(f"Error type: {type(e)}")
+        import traceback
+        traceback.print_exc()
         return f"An error occurred: {str(e)}"
 
 # Create Gradio interface
@@ -280,4 +307,5 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch() 
+    print("Starting Gradio interface...")
+    iface.launch(debug=True)