Update app.py

app.py

@@ -1,31 +1,419 @@
 import gradio as gr
-import random
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import re
 
-# Simple fallback responses
-fallback_responses = [
-    "I understand you want me to explain quantum computing. It's a complex topic that deals with how very small particles behave differently from objects in our everyday world.",
-    "Quantum computing uses quantum bits or qubits, which can exist in multiple states at once, unlike regular computer bits that are only 0 or 1.",
-    "This allows quantum computers to solve certain problems much faster than traditional computers.",
-    "I'd be happy to help you learn about quantum computing! It's a fascinating field of physics and computer science.",
-    "Quantum computers use phenomena like superposition and entanglement to perform calculations in new ways."
-]
+# Initialize the model and tokenizer
+@torch.no_grad()
+def load_model():
+    print("Loading Qwen3-0.6B model...")
+    
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            "Qwen/Qwen3-0.6B",
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        print("Qwen3-0.6B model loaded successfully!")
+        return tokenizer, model
+        
+    except Exception as e:
+        print(f"Error loading Qwen3-0.6B: {e}")
+        return None, None
+
+# Load the model
+tokenizer, model = load_model()
+
+def remove_think_tags(text):
+    """
+    Remove <think>...</think> tags from text - METHOD 1
+    """
+    cleaned_text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
+    return cleaned_text.strip()
+
+def generate_response(message, history, temperature=0.7, max_length=256):
+    """
+    Generate a response using Qwen3-0.6B with your specified method
+    """
+    if tokenizer is None or model is None:
+        return "⚠️ Model is not loaded properly. Please check the console logs."
+    
+    try:
+        # Convert history to messages format
+        messages = []
+        for human_msg, assistant_msg in history:
+            messages.extend([
+                {"role": "user", "content": human_msg},
+                {"role": "assistant", "content": assistant_msg}
+            ])
+        
+        # Add current message
+        messages.append({"role": "user", "content": message})
+        
+        # Apply chat template exactly as in your example
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(model.device)
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_length,
+                temperature=temperature,
+                do_sample=True if temperature > 0.1 else False,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        # Extract only the new generated text
+        response = tokenizer.decode(
+            outputs[0][inputs["input_ids"].shape[-1]:], 
+            skip_special_tokens=True
+        )
+        
+        # Clean think tags from response
+        clean_response = remove_think_tags(response)
+        
+        return clean_response if clean_response else "I'm here to help! What would you like to know?"
+        
+    except Exception as e:
+        print(f"Generation error: {e}")
+        return f"⚠️ Error generating response: {str(e)}"
 
-def simple_chat(message, history):
-    if not message.strip():
-        return "", history
+def chat_interface(message, history, temperature, max_length):
+    """
+    Main chat interface function
+    """
+    if not message or not message.strip():
+        return "", history or []
     
-    # Simulate AI thinking
-    response = random.choice(fallback_responses)
+    # Generate response
+    bot_response = generate_response(message, history or [], temperature, max_length)
     
     # Update history
-    new_history = history + [[message, response]]
+    new_history = (history or []) + [[message, bot_response]]
     
     return "", new_history
 
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(label="Your message")
+def clear_chat():
+    """
+    Clear the chat history
+    """
+    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
+    new_history = last_conversation + [[last_user_message, bot_response]]
+    
+    return new_history
+
+# Custom CSS for beautiful UI
+custom_css = """
+.gradio-container {
+    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    min-height: 100vh;
+    padding: 20px;
+}
+
+.main-container {
+    max-width: 1200px;
+    margin: 0 auto;
+    background: white;
+    border-radius: 20px;
+    box-shadow: 0 20px 40px rgba(0,0,0,0.1);
+    overflow: hidden;
+}
+
+.header {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    color: white;
+    padding: 30px;
+    text-align: center;
+}
+
+.header h1 {
+    margin: 0;
+    font-size: 2.5em;
+    font-weight: 700;
+}
+
+.header p {
+    margin: 10px 0 0 0;
+    opacity: 0.9;
+    font-size: 1.2em;
+}
+
+.content {
+    display: flex;
+    min-height: 600px;
+}
+
+.chat-column {
+    flex: 3;
+    display: flex;
+    flex-direction: column;
+}
+
+.control-column {
+    flex: 1;
+    background: #f8f9fa;
+    padding: 25px;
+    border-left: 1px solid #e1e5e9;
+}
+
+.chatbot-container {
+    flex: 1;
+    display: flex;
+    flex-direction: column;
+    min-height: 500px;
+}
+
+#chatbot {
+    flex: 1;
+    min-height: 500px !important;
+    border: none !important;
+    background: white !important;
+    padding: 20px !important;
+    margin: 0 !important;
+}
+
+#chatbot .message {
+    padding: 15px 20px !important;
+    margin: 10px 0 !important;
+    border-radius: 15px !important;
+    max-width: 80% !important;
+}
+
+#chatbot .user-message {
+    background: #667eea !important;
+    color: white !important;
+    margin-left: auto !important;
+}
+
+#chatbot .bot-message {
+    background: #f1f3f4 !important;
+    color: #333 !important;
+    margin-right: auto !important;
+}
+
+.input-container {
+    background: #f8f9fa;
+    padding: 20px;
+    border-top: 1px solid #e1e5e9;
+}
+
+.control-panel {
+    background: white;
+    padding: 20px;
+    border-radius: 15px;
+    margin-bottom: 20px;
+    box-shadow: 0 4px 12px rgba(0,0,0,0.1);
+}
+
+.model-info {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    color: white;
+    padding: 20px;
+    border-radius: 15px;
+}
+
+.gr-button {
+    background: linear-gradient(45deg, #667eea, #764ba2) !important;
+    border: none !important;
+    color: white !important;
+    border-radius: 10px !important;
+    padding: 12px 24px !important;
+    font-weight: 600 !important;
+    margin: 5px !important;
+}
+
+.clear-btn {
+    background: linear-gradient(45deg, #ff6b6b, #ee5a24) !important;
+}
+
+.retry-btn {
+    background: linear-gradient(45deg, #00b894, #00a085) !important;
+}
+
+.textbox {
+    border-radius: 12px !important;
+    border: 2px solid #e1e5e9 !important;
+    padding: 15px !important;
+    font-size: 16px !important;
+}
+
+.textbox:focus {
+    border-color: #667eea !important;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+}
+
+.examples-panel {
+    background: white;
+    padding: 20px;
+    border-top: 1px solid #e1e5e9;
+}
+
+.loading {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    height: 200px;
+    color: #666;
+}
+"""
+
+# Create the Gradio interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    
+    with gr.Column(elem_classes="main-container"):
+        # Header
+        with gr.Column(elem_classes="header"):
+            gr.Markdown("# 🤖 Qwen3-0.6B Chatbot")
+            gr.Markdown("Chat with Alibaba's advanced Qwen3-0.6B model - Think tags automatically removed!")
+        
+        with gr.Row(elem_classes="content"):
+            # Left Column - Chat (70%)
+            with gr.Column(elem_classes="chat-column"):
+                with gr.Column(elem_classes="chatbot-container"):
+                    chatbot = gr.Chatbot(
+                        value=[],
+                        label="",
+                        elem_id="chatbot",
+                        show_copy_button=True,
+                        avatar_images=("👤", "🤖"),
+                        height=500,
+                        container=True,
+                        show_label=False
+                    )
+                
+                with gr.Column(elem_classes="input-container"):
+                    with gr.Row():
+                        msg = gr.Textbox(
+                            label="",
+                            placeholder="Type your message here...",
+                            lines=2,
+                            scale=4,
+                            container=False,
+                            show_label=False
+                        )
+                        submit_btn = gr.Button("Send 🚀", size="lg", scale=1)
+                    
+                    with gr.Row():
+                        clear_btn = gr.Button("🗑️ Clear Chat", elem_classes="clear-btn")
+                        retry_btn = gr.Button("🔄 Retry Last", elem_classes="retry-btn")
+                        gr.HTML("""<div style="flex: 1; text-align: center; color: #666; font-size: 12px; padding: 10px;">
+                            Press Enter to send • Shift+Enter for new line
+                        </div>""")
+            
+            # Right Column - Controls (30%)
+            with gr.Column(elem_classes="control-column"):
+                with gr.Column(elem_classes="control-panel"):
+                    gr.Markdown("### ⚙️ Generation Settings")
+                    
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.5,
+                        value=0.7,
+                        step=0.1,
+                        label="Temperature",
+                        info="Lower = more predictable, Higher = more creative"
+                    )
+                    
+                    max_length = gr.Slider(
+                        minimum=50,
+                        maximum=1000,
+                        value=256,
+                        step=50,
+                        label="Max Response Length",
+                        info="Tokens in generated response"
+                    )
+                
+                with gr.Column(elem_classes="model-info"):
+                    gr.Markdown("### ℹ️ Model Info")
+                    if tokenizer and model:
+                        gr.Markdown("""
+                        **Model:** Qwen3-0.6B ✅
+                        **Status:** Ready to chat!
+                        **Think Tags:** Auto-removed ✅
+                        
+                        **Features:**
+                        • 0.6B parameters
+                        • 128K context
+                        • Multilingual
+                        """)
+                    else:
+                        gr.Markdown("""
+                        **Status:** ⚠️ Loading failed
+                        **Check console for errors**
+                        """)
+        
+        # Examples Section
+        with gr.Column(elem_classes="examples-panel"):
+            gr.Markdown("### 💡 Try These Examples")
+            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 exploration",
+                    "Help me plan a healthy weekly meal plan"
+                ],
+                inputs=msg,
+                label="Click any example to start chatting!",
+                examples_per_page=6
+            )
+
+    # Event handlers
+    submit_event = msg.submit(
+        chat_interface,
+        inputs=[msg, chatbot, temperature, max_length],
+        outputs=[msg, chatbot]
+    )
+    
+    submit_btn.click(
+        chat_interface,
+        inputs=[msg, chatbot, temperature, max_length],
+        outputs=[msg, chatbot]
+    )
+    
+    clear_btn.click(
+        clear_chat,
+        outputs=[chatbot]
+    )
     
-    msg.submit(simple_chat, [msg, chatbot], [msg, chatbot])
+    retry_btn.click(
+        retry_last_response,
+        inputs=[chatbot, temperature, max_length],
+        outputs=[chatbot]
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        share=False,
+        show_error=True
+    )