Update app.py

app.py

@@ -12,91 +12,49 @@ model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
 
 print("Model loaded successfully!")
 
-# Define the prediction function that works with the modern format
+# Define the prediction function that works with the modern 'messages' format
 def predict(message, history):
-    try:
-        # Validate input
-        if not message or not message.strip():
-            return "Please enter a message."
-        
-        # Format history for DialoGPT - handle both old and new formats
-        history_transformer_format = ""
-        
-        # Handle the new 'messages' format (list of dicts)
-        if history and isinstance(history[0], dict):
-            for turn in history:
-                if turn.get("role") == "user":
-                    history_transformer_format += turn["content"] + tokenizer.eos_token
-                elif turn.get("role") == "assistant":
-                    history_transformer_format += turn["content"] + tokenizer.eos_token
-        
-        # Handle the old 'tuples' format (list of lists)
-        elif history and isinstance(history[0], list):
-            for exchange in history:
-                if len(exchange) >= 2:
-                    user_msg, bot_msg = exchange[0], exchange[1]
-                    if user_msg:
-                        history_transformer_format += str(user_msg) + tokenizer.eos_token
-                    if bot_msg:
-                        history_transformer_format += str(bot_msg) + tokenizer.eos_token
-        
-        # Add the current message
-        input_text = history_transformer_format + str(message) + tokenizer.eos_token
-
-        # Tokenize the input
-        new_user_input_ids = tokenizer.encode(input_text, return_tensors='pt')
-
-        # Generate a response with memory management
-        with torch.no_grad():
-            bot_output_ids = model.generate(
-                new_user_input_ids,
-                max_length=1000,  # Reduced for better performance
-                pad_token_id=tokenizer.eos_token_id,
-                no_repeat_ngram_size=3,
-                do_sample=True,
-                top_k=50,
-                top_p=0.7,
-                temperature=0.8
-            )
+    # Format the history for DialoGPT. It expects a flat string of alternating user/bot messages.
+    history_transformer_format = ""
+    for user_msg, bot_msg in history:
+        history_transformer_format += user_msg + tokenizer.eos_token
+        history_transformer_format += bot_msg + tokenizer.eos_token
+    
+    # Append the new user message
+    history_transformer_format += message + tokenizer.eos_token
+
+    # Tokenize the input
+    new_user_input_ids = tokenizer.encode(history_transformer_format, return_tensors='pt')
+
+    # Generate a response
+    bot_output_ids = model.generate(
+        new_user_input_ids,
+        max_length=1250,
+        pad_token_id=tokenizer.eos_token_id,
+        no_repeat_ngram_size=3,
+        do_sample=True,
+        top_k=100,
+        top_p=0.7,
+        temperature=0.8
+    )
 
-        # Decode the response
-        response = tokenizer.decode(
-            bot_output_ids[:, new_user_input_ids.shape[-1]:][0], 
-            skip_special_tokens=True
-        ).strip()
-        
-        # Clean up and validate response
-        if not response:
-            response = "I'm sorry, I couldn't generate a response. Could you try rephrasing your question?"
-        
-        # Limit response length to prevent protocol errors
-        if len(response) > 500:
-            response = response[:500] + "..."
-        
-        return response
-        
-    except Exception as e:
-        print(f"Error in predict function: {str(e)}")
-        return "Sorry, I encountered an error. Please try again with a different message."
+    # Decode the response, skipping the input part
+    response = tokenizer.decode(bot_output_ids[:, new_user_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    
+    return response
 
-# Create a simple ChatInterface
+# Build the Gradio interface using the modern 'gr.ChatInterface'
+# This is much simpler and handles all the UI elements for you.
 demo = gr.ChatInterface(
     fn=predict,
     title="DialoGPT-medium Chatbot",
-    description="Chat with Microsoft's DialoGPT-medium model!",
-    examples=[
-        "Hello!",
-        "How are you?",
-        "Tell me a joke",
-        "What's the weather like?"
-    ],
-    cache_examples=False
+    description="This chatbot uses the microsoft/DialoGPT-medium model. Start typing to chat!",
+    theme="soft",
+    examples=["Hello!", "How does a computer work?", "Tell me a joke."],
+    undo_btn="Undo Last Turn",
+    clear_btn="Clear Chat",
 )
 
-# Launch the app with public sharing enabled
+# Launch the app. No 'share=True' is needed on Spaces.
 if __name__ == "__main__":
-    demo.launch(
-        share=True,  # This creates the public link
-        server_name="0.0.0.0",
-        server_port=7860
-    )
+    demo.queue().launch()