Update app.py

app.py

@@ -6,56 +6,89 @@ import torch
 model_id = "jatingocodeo/SmolLM2"
 
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-    return model, tokenizer
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        # Ensure the tokenizer has the necessary special tokens
+        special_tokens = {
+            'pad_token': '[PAD]',
+            'eos_token': '</s>',
+            'bos_token': '<s>'
+        }
+        tokenizer.add_special_tokens(special_tokens)
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            pad_token_id=tokenizer.pad_token_id
+        )
+        # Resize token embeddings to match new tokenizer
+        model.resize_token_embeddings(len(tokenizer))
+        return model, tokenizer
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        raise
 
 def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
-    # Load model and tokenizer (caching them for subsequent calls)
-    if not hasattr(generate_text, "model"):
-        generate_text.model, generate_text.tokenizer = load_model()
-    
-    # Encode the prompt
-    input_ids = generate_text.tokenizer.encode(prompt, return_tensors="pt")
-    input_ids = input_ids.to(generate_text.model.device)
+    try:
+        # Load model and tokenizer (caching them for subsequent calls)
+        if not hasattr(generate_text, "model"):
+            generate_text.model, generate_text.tokenizer = load_model()
+        
+        # Ensure the prompt is not empty
+        if not prompt.strip():
+            return "Please enter a prompt."
+        
+        # Add BOS token if needed
+        if not prompt.startswith(generate_text.tokenizer.bos_token):
+            prompt = generate_text.tokenizer.bos_token + 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)
+        
+        # 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
+                temperature=temperature,
+                top_k=top_k,
+                do_sample=True,
+                pad_token_id=generate_text.tokenizer.pad_token_id,
+                eos_token_id=generate_text.tokenizer.eos_token_id,
+                num_return_sequences=1
+            )
+        
+        # Decode and return the generated text
+        generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        return generated_text.strip()
     
-    # Generate text
-    with torch.no_grad():
-        output_ids = generate_text.model.generate(
-            input_ids,
-            max_length=max_length,
-            temperature=temperature,
-            top_k=top_k,
-            pad_token_id=generate_text.tokenizer.pad_token_id,
-            eos_token_id=generate_text.tokenizer.eos_token_id,
-            do_sample=True
-        )
-    
-    # Decode and return the generated text
-    generated_text = generate_text.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return generated_text
+    except Exception as e:
+        print(f"Error during generation: {str(e)}")
+        return f"An error occurred: {str(e)}"
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt here...", lines=2),
         gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K"),
     ],
-    outputs=gr.Textbox(label="Generated Text"),
+    outputs=gr.Textbox(label="Generated Text", lines=5),
     title="SmolLM2 Text Generator",
-    description="Generate text using the fine-tuned SmolLM2 model",
+    description="""Generate text using the fine-tuned SmolLM2 model.
+    - Max Length: Controls the length of generated text
+    - Temperature: Controls randomness (higher = more creative)
+    - Top K: Controls diversity of word choices""",
     examples=[
         ["Once upon a time", 100, 0.7, 50],
         ["The quick brown fox", 150, 0.8, 40],
         ["In a galaxy far far away", 200, 0.9, 30],
-    ]
+    ],
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":