Update app.py

app.py

@@ -1,63 +1,83 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 import torch
 import time
 
-# Load the model and tokenizer
+# Initialize the pipeline
 @torch.no_grad()
 def load_model():
-    print("Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
-    model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
-    
-    # Add padding token if it doesn't exist
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
-    
+    print("Loading Qwen3-0.6B model...")
+    pipe = pipeline(
+        "text-generation",
+        model="Qwen/Qwen3-0.6B",
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
     print("Model loaded successfully!")
-    return tokenizer, model
+    return pipe
 
-# Initialize the model
-tokenizer, model = load_model()
+# Load the model
+pipe = load_model()
 
-def generate_response(message, history, temperature=0.7, max_length=150):
+def format_chat_template(messages):
     """
-    Generate a response using GPT-2
+    Format messages using the model's chat template
     """
     try:
-        # Format the conversation history
-        conversation_history = ""
-        for human, assistant in history:
-            conversation_history += f"Human: {human}\nAssistant: {assistant}\n"
-        
-        # Add the current message
-        full_prompt = f"{conversation_history}Human: {message}\nAssistant:"
+        # Use the model's built-in chat template
+        formatted_prompt = pipe.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        return formatted_prompt
+    except Exception as e:
+        # Fallback formatting
+        conversation = ""
+        for msg in messages:
+            if msg["role"] == "user":
+                conversation += f"User: {msg['content']}\n\nAssistant:"
+            elif msg["role"] == "assistant":
+                conversation += f" {msg['content']}\n\n"
+        return conversation
+
+def generate_response(message, history, temperature=0.7, max_length=512):
+    """
+    Generate a response using Qwen3-0.6B
+    """
+    try:
+        # Convert Gradio history to messages format
+        messages = []
+        for human_msg, assistant_msg in history:
+            messages.extend([
+                {"role": "user", "content": human_msg},
+                {"role": "assistant", "content": assistant_msg}
+            ])
         
-        # Tokenize the input
-        inputs = tokenizer.encode(full_prompt, return_tensors="pt", max_length=1024, truncation=True)
+        # Add current message
+        messages.append({"role": "user", "content": message})
         
         # Generate response
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs,
-                max_length=inputs.shape[1] + max_length,
-                temperature=temperature,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id,
-                top_p=0.9,
-                repetition_penalty=1.1,
-                early_stopping=True
-            )
+        formatted_prompt = format_chat_template(messages)
         
-        # Decode the response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        outputs = pipe(
+            formatted_prompt,
+            max_new_tokens=max_length,
+            temperature=temperature,
+            do_sample=True,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            pad_token_id=pipe.tokenizer.eos_token_id,
+            eos_token_id=pipe.tokenizer.eos_token_id,
+            return_full_text=False  # Only return the generated part
+        )
         
-        # Extract only the new generated text
-        response = response[len(full_prompt):].strip()
+        response = outputs[0]['generated_text'].strip()
         
-        # Clean up the response (remove any trailing human prompts)
-        if "Human:" in response:
-            response = response.split("Human:")[0].strip()
+        # Clean up response
+        if "User:" in response:
+            response = response.split("User:")[0].strip()
         
         return response
         
@@ -85,52 +105,115 @@ def clear_chat():
     """
     return []
 
+def retry_last_response(history, temperature, max_length):
+    """
+    Retry the last user message
+    """
+    if not history:
+        return history
+    
+    # Remove the last assistant response
+    last_conversation = history[:-1]
+    last_user_message = history[-1][0]
+    
+    # Regenerate response
+    bot_response = generate_response(last_user_message, last_conversation, temperature, max_length)
+    
+    # Update history
+    last_conversation.append([last_user_message, bot_response])
+    
+    return last_conversation
+
 # Custom CSS for beautiful styling
 custom_css = """
-#chatbot {
-    background-color: #f8f9fa;
-    border-radius: 10px;
-    padding: 20px;
-    height: 500px;
-}
-
 .gradio-container {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    min-height: 100vh;
 }
 
 .contain {
     background-color: white;
-    border-radius: 15px;
-    padding: 20px;
-    box-shadow: 0 10px 30px rgba(0,0,0,0.1);
+    border-radius: 20px;
+    padding: 30px;
+    box-shadow: 0 20px 40px rgba(0,0,0,0.1);
+    margin: 20px;
 }
 
 .dark .contain {
     background-color: #1e1e1e;
 }
 
+#chatbot {
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
+    border-radius: 15px;
+    padding: 25px;
+    height: 500px;
+    border: none;
+    box-shadow: inset 0 2px 10px rgba(0,0,0,0.1);
+}
+
+.dark #chatbot {
+    background: linear-gradient(135deg, #2d3748 0%, #4a5568 100%);
+}
+
 .gr-button {
     background: linear-gradient(45deg, #FF6B6B, #FF8E53) !important;
     border: none !important;
     color: white !important;
     border-radius: 25px !important;
-    padding: 10px 20px !important;
+    padding: 12px 25px !important;
+    font-weight: 600 !important;
+    transition: all 0.3s ease !important;
 }
 
 .gr-button:hover {
     transform: translateY(-2px);
-    box-shadow: 0 5px 15px rgba(255,107,107,0.4);
+    box-shadow: 0 8px 20px rgba(255,107,107,0.4);
 }
 
-.slider-container {
+.control-panel {
     background: white;
-    padding: 15px;
-    border-radius: 10px;
-    margin: 10px 0;
+    padding: 25px;
+    border-radius: 15px;
+    box-shadow: 0 5px 15px rgba(0,0,0,0.1);
+    border: 1px solid #e1e5e9;
 }
 
-.dark .slider-container {
+.dark .control-panel {
     background: #2d2d2d;
+    border-color: #404040;
+}
+
+.slider-container {
+    margin: 15px 0;
+}
+
+.model-info {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    color: white;
+    padding: 20px;
+    border-radius: 15px;
+    margin: 15px 0;
+}
+
+.textbox {
+    border-radius: 15px !important;
+    padding: 15px !important;
+    border: 2px solid #e1e5e9 !important;
+}
+
+.textbox:focus {
+    border-color: #667eea !important;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+}
+
+.markdown-container h1 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    text-fill-color: transparent;
+    font-weight: 700 !important;
 }
 """
 
@@ -138,63 +221,76 @@ custom_css = """
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🤖 GPT-2 Chatbot
-        **Chat with OpenAI's GPT-2 model!** This AI assistant can help you with conversations, creative writing, and more.
-        """
+        # 🚀 Qwen3-0.6B Chatbot
+        **Chat with Alibaba's advanced Qwen3-0.6B model!** Experience more fluent and intelligent conversations.
+        """,
+        elem_classes="markdown-container"
     )
     
-    with gr.Row():
+    with gr.Row(equal_height=False):
         with gr.Column(scale=3):
             chatbot = gr.Chatbot(
-                label="Chat History",
+                label="💬 Chat with Qwen3",
                 elem_id="chatbot",
                 show_copy_button=True,
-                avatar_images=("👤", "🤖")
+                avatar_images=("👤", "🤖"),
+                height=500,
+                bubble_full_width=False
             )
             
             with gr.Row():
                 msg = gr.Textbox(
-                    label="Type your message here...",
-                    placeholder="Hello! How can I help you today?",
+                    label="💭 Your message",
+                    placeholder="Ask me anything...",
                     lines=2,
-                    scale=4
+                    scale=4,
+                    container=False
                 )
-                submit_btn = gr.Button("Send 🚀", scale=1)
+                with gr.Column(scale=1):
+                    submit_btn = gr.Button("Send 🚀", size="lg")
             
             with gr.Row():
-                clear_btn = gr.Button("Clear Chat 🗑️")
+                clear_btn = gr.Button("🗑️ Clear Chat", size="sm")
+                retry_btn = gr.Button("🔄 Retry", size="sm")
+                gr.HTML("""<div style="text-align: center; color: #666; font-size: 12px;">
+                    Pro tip: Use Shift+Enter for new line, Enter to send
+                </div>""")
         
         with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Settings")
-            
-            with gr.Group(elem_classes="slider-container"):
+            with gr.Group(elem_classes="control-panel"):
+                gr.Markdown("### ⚙️ Generation Settings")
+                
                 temperature = gr.Slider(
                     minimum=0.1,
-                    maximum=1.0,
+                    maximum=1.5,
                     value=0.7,
                     step=0.1,
-                    label="Temperature",
-                    info="Higher = more creative, Lower = more focused"
+                    label="🎛️ Temperature",
+                    info="Lower = more deterministic, Higher = more creative"
                 )
                 
                 max_length = gr.Slider(
-                    minimum=50,
-                    maximum=300,
-                    value=150,
-                    step=10,
-                    label="Max Response Length",
-                    info="Maximum length of generated responses"
+                    minimum=64,
+                    maximum=1024,
+                    value=512,
+                    step=64,
+                    label="📏 Max Response Length",
+                    info="Tokens in generated response"
                 )
             
-            gr.Markdown("### ℹ️ About")
-            gr.Markdown("""
-            This chatbot uses **GPT-2** from OpenAI.
-            
-            **Tips:**
-            - Be specific in your questions
-            - Adjust temperature for creativity
-            - Clear chat to start fresh
-            """)
+            with gr.Group(elem_classes="model-info"):
+                gr.Markdown("### ℹ️ Model Info")
+                gr.Markdown("""
+                **Model:** Qwen3-0.6B  
+                **Provider:** Alibaba  
+                **Context:** 128K tokens  
+                **Language:** Multilingual  
+                
+                💡 **Tips:**
+                - Be specific and clear
+                - Use proper punctuation
+                - Adjust temperature for creativity
+                """)
     
     # Event handlers
     submit_event = msg.submit(
@@ -214,14 +310,29 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         outputs=[chatbot]
     )
     
-    # Additional handlers for Enter key
-    submit_event.then(
-        lambda: gr.Textbox(value="", interactive=True),
-        outputs=[msg]
+    retry_btn.click(
+        retry_last_response,
+        inputs=[chatbot, temperature, max_length],
+        outputs=[chatbot]
     )
+    
+    # Additional examples
+    with gr.Accordion("💡 Example Prompts", open=False):
+        gr.Examples(
+            examples=[
+                "Explain quantum computing in simple terms",
+                "Write a short poem about artificial intelligence",
+                "What are the benefits of renewable energy?",
+                "How do I learn programming effectively?",
+                "Tell me an interesting fact about space"
+            ],
+            inputs=msg,
+            label="Click any example to try it out!"
+        )
 
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
-        share=False
+        share=False,
+        show_error=True
     )