Update app.py

app.py

@@ -2,53 +2,61 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Load your model and tokenizer
+# Model name
 MODEL_NAME = "Dhansh2001/my-fitlien-chatbot-pruned-quantized"
 
+# Load model & tokenizer
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-    
+
+    # Ensure pad token is set
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
         model.config.pad_token_id = tokenizer.eos_token_id
-    
+
     return tokenizer, model
 
-# Load model once at startup
+# Load once at startup
 tokenizer, model = load_model()
 
+# Use CPU or GPU automatically
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
 def chat_with_bot(message, history):
-    # Encode the new user input, add the eos_token and return a tensor in Pytorch
-    new_user_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
-    
-    # Create attention mask
-    attention_mask = torch.ones(new_user_input_ids.shape, dtype=torch.long)
-    
-    # Generate a response while limiting the total chat history to 1000 tokens
-    with torch.no_grad():
-        chat_history_ids = model.generate(
-            new_user_input_ids, 
-            attention_mask=attention_mask,
-            max_length=1000,
-            num_beams=5,
-            no_repeat_ngram_size=3,
-            do_sample=True,
-            temperature=0.7,
-            pad_token_id=tokenizer.eos_token_id,
-            early_stopping=True
-        )
-    
-    # Decode the response
-    response = tokenizer.decode(chat_history_ids[:, new_user_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    
-    return response
+    try:
+        # Encode user input
+        inputs = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt").to(device)
+
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_length=500,
+                num_beams=5,
+                no_repeat_ngram_size=3,
+                do_sample=True,
+                temperature=0.7,
+                pad_token_id=tokenizer.eos_token_id,
+                early_stopping=True
+            )
+
+        # Decode and return
+        reply = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
+
+        # Ensure a safe string return
+        return reply if reply.strip() != "" else "I'm not sure how to answer that."
+
+    except Exception as e:
+        # Catch errors so API never 500s silently
+        return f"⚠️ Error: {str(e)}"
 
-# Create Gradio interface
+# Gradio ChatInterface
 iface = gr.ChatInterface(
     fn=chat_with_bot,
     title="My Fitlien Chatbot",
-    description="A chatbot fine-tuned from DialoGPT-medium for fitness conversations",
+    description="A chatbot fine-tuned from DialoGPT-medium for fitness conversations.",
     examples=[
         "Hello! How are you?",
         "What's a good workout routine?",
@@ -61,4 +69,4 @@ iface = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()