Update app.py

app.py

@@ -1,33 +1,32 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-import gradio as gr
-
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
-
-chat_history_ids = None
-
 def chat(user_input, history=[]):
     global chat_history_ids
 
-    # Encode input + add to chat history
+    # Encode user input
     new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
+
+    # Append to chat history or start new
     bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1) if chat_history_ids is not None else new_input_ids
 
-    # Generate reply
-    chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    # Generate response (LIMIT max_length and num_return_sequences)
+    chat_history_ids = model.generate(
+        bot_input_ids,
+        max_length=1000,           # Can reduce to 500 if needed
+        pad_token_id=tokenizer.eos_token_id,
+        do_sample=True,
+        top_k=50,
+        top_p=0.95,
+        temperature=0.7,
+        num_return_sequences=1
+    )
+
+    # Decode only the new response part
     response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
 
+    # Truncate response if it’s too long (hard limit)
+    if len(response) > 1000:
+        response = response[:1000] + "..."
+
     history.append((user_input, response))
     return history, history
 
-# Chat UI
-chatbot_ui = gr.ChatInterface(
-    fn=chat,
-    title="Teen Mental Health Chatbot 🤖💬",
-    description="Talk to a supportive AI. Not a replacement for professional help.",
-)
-
-chatbot_ui.launch()