Update app.py

app.py

@@ -12,81 +12,206 @@ if token:
     login(token=token)
     print("Successfully logged in to Hugging Face!")
 
-# Load the base model and tokenizer with better error handling
-model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+# Use a more reliable model that works well in Spaces
+model_name = "microsoft/DialoGPT-medium"
 print("Loading tokenizer...")
 
 try:
-    # Try loading with use_fast=False to avoid tokenizer issues
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name, 
-        use_fast=False,
-        trust_remote_code=True
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    print("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto"
     )
+    print("Model loaded successfully!")
+    
 except Exception as e:
-    print(f"Error loading tokenizer: {e}")
-    print("Trying alternative tokenizer...")
-    # Fallback to a known working tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(
-        "huggingface/CodeBERTa-small-v1",
-        use_fast=False
-    )
+    print(f"Error loading model: {e}")
+    print("Using a simpler fallback...")
+    # Fallback to an even simpler model
+    model_name = "gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    print("Fallback model loaded!")
 
 # Set pad token if not set
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-print("Loading model...")
-# Load base model with better memory management
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    offload_folder="./offload"  # Add offload directory
-)
-
-print("Model loaded successfully!")
-
-def chat_function(message, history):
+def chat_function(message, temperature, max_length):
+    if not message or not message.strip():
+        return "Please enter a message to get started! 💬"
+    
+    # Limit input length
+    if len(message) > 500:
+        return "⚠️ Message too long! Please keep it under 500 characters."
+    
     try:
-        # Simple format for Mistral
-        prompt = f"<s>[INST] {message} [/INST]"
+        # Encode the input with length limit
+        message = message.strip()[:200]  # Hard limit on input
+        input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
         
-        # Tokenize
-        inputs = tokenizer(
-            prompt, 
-            return_tensors="pt", 
-            truncation=True, 
-            max_length=1024
-        )
+        # Ensure input isn't too long
+        if input_ids.shape[-1] > 100:
+            input_ids = input_ids[:, :100]
         
-        # Generate
+        # Generate response with controlled parameters
         with torch.no_grad():
-            outputs = model.generate(
-                inputs.input_ids,
-                max_new_tokens=256,
-                temperature=0.7,
+            chat_history_ids = model.generate(
+                input_ids, 
+                max_length=min(input_ids.shape[-1] + max_length, 200),  # Total max 200 tokens
+                temperature=temperature,
                 do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=tokenizer.eos_token_id,
+                num_return_sequences=1,
+                repetition_penalty=1.1,
+                top_p=0.9
             )
         
         # Decode response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response = response.replace(prompt, "").strip()
+        response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
         
+        # Clean up the response
+        response = response.strip()
+        if not response:
+            responses = [
+                "I'm here to help! What would you like to know? 🤔",
+                "How can I assist you today? 😊",
+                "What's on your mind? 💭",
+                "I'm ready to chat! What would you like to discuss? 🗣️"
+            ]
+            import random
+            response = random.choice(responses)
+            
         return response
         
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"🚫 Sorry, I encountered an error: {str(e)}"
+
+def clear_chat():
+    return "", ""
+
+# Custom CSS for better styling
+css = """
+.gradio-container {
+    max-width: 800px !important;
+    margin: auto !important;
+}
+.chat-container {
+    border-radius: 15px;
+    box-shadow: 0 4px 15px rgba(0,0,0,0.1);
+}
+.input-container {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    padding: 20px;
+    border-radius: 15px;
+    margin-bottom: 10px;
+}
+.output-container {
+    background: #f8f9fa;
+    border-radius: 15px;
+    padding: 20px;
+}
+"""
 
-# Create the simplest possible interface
-demo = gr.Interface(
-    fn=chat_function,
-    inputs=gr.Textbox(placeholder="Type your message here..."),
-    outputs=gr.Textbox(),
-    title="3AI Custom Mistral Model",
-    description="Chat with the Mistral 7B model"
-)
+# Create a refined interface
+with gr.Blocks(title="3AI Chat Bot", css=css, theme=gr.themes.Soft()) as demo:
+    gr.HTML("""
+    <div style="text-align: center; padding: 20px;">
+        <h1 style="color: #2c3e50; margin-bottom: 10px;">🤖 3AI Chat Bot</h1>
+        <p style="color: #7f8c8d; font-size: 16px;">Powered by AI • Fast • Intelligent • Helpful</p>
+    </div>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Group(elem_classes="input-container"):
+                gr.Markdown("### 💬 **Chat with AI**")
+                message_input = gr.Textbox(
+                    placeholder="Type your message here... (max 500 characters)", 
+                    label="Your Message",
+                    lines=3,
+                    max_lines=5,
+                    show_label=False,
+                    container=False
+                )
+                
+                with gr.Row():
+                    submit_btn = gr.Button("🚀 Send", variant="primary", scale=2)
+                    clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
+        
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown("### ⚙️ **Settings**")
+                temperature = gr.Slider(
+                    minimum=0.1, 
+                    maximum=1.0, 
+                    value=0.7, 
+                    step=0.1,
+                    label="🌡️ Creativity",
+                    info="Higher = more creative"
+                )
+                max_length = gr.Slider(
+                    minimum=20, 
+                    maximum=100, 
+                    value=50, 
+                    step=10,
+                    label="📏 Response Length",
+                    info="Max tokens to generate"
+                )
+    
+    with gr.Group(elem_classes="output-container"):
+        gr.Markdown("### 🤖 **AI Response**")
+        response_output = gr.Textbox(
+            label="",
+            lines=6,
+            max_lines=10,
+            interactive=False,
+            show_label=False,
+            container=False,
+            placeholder="AI responses will appear here... ✨"
+        )
+    
+    # Character counter
+    char_count = gr.HTML("<p style='text-align: right; color: #7f8c8d; font-size: 12px;'>0/500 characters</p>")
+    
+    # Event handlers
+    submit_btn.click(
+        fn=chat_function,
+        inputs=[message_input, temperature, max_length],
+        outputs=response_output
+    )
+    
+    message_input.submit(
+        fn=chat_function,
+        inputs=[message_input, temperature, max_length],
+        outputs=response_output
+    )
+    
+    clear_btn.click(
+        fn=clear_chat,
+        outputs=[message_input, response_output]
+    )
+    
+    # Update character counter
+    def update_char_count(text):
+        count = len(text) if text else 0
+        color = "#e74c3c" if count > 500 else "#7f8c8d"
+        return f"<p style='text-align: right; color: {color}; font-size: 12px;'>{count}/500 characters</p>"
+    
+    message_input.change(
+        fn=update_char_count,
+        inputs=message_input,
+        outputs=char_count
+    )
+    
+    # Add footer
+    gr.HTML("""
+    <div style="text-align: center; padding: 20px; color: #7f8c8d; font-size: 12px;">
+        <p>Made with ❤️ using Gradio • Powered by Transformers</p>
+    </div>
+    """)
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()