Upload app.py with huggingface_hub

app.py

@@ -1,88 +1,116 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Initialize the sentiment analysis pipeline
-sentiment_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
+# Model: Mistral 7B Instruct - powerful open model
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
+
+print(f"Loading model: {MODEL_NAME}...")
 
 # Initialize text generation pipeline
-generator = pipeline("text-generation", model="distilgpt2", max_length=100)
+try:
+    generator = pipeline(
+        "text-generation",
+        model=MODEL_NAME,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        trust_remote_code=True
+    )
+    print("Model loaded successfully!")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    # Fallback to a smaller model
+    MODEL_NAME = "microsoft/DialoGPT-medium"
+    generator = pipeline("text-generation", model=MODEL_NAME)
+    print(f"Loaded fallback model: {MODEL_NAME}")
 
-def analyze_sentiment(text):
-    """Analyze the sentiment of the input text."""
-    if not text.strip():
-        return "Please enter some text to analyze."
+def chat_with_ai(message, history):
+    """Chat with the AI model."""
+    if not message.strip():
+        return history
     
-    result = sentiment_pipeline(text)[0]
-    label = result['label']
-    score = result['score']
+    # Build conversation prompt
+    conversation = ""
+    for user_msg, assistant_msg in history:
+        conversation += f"<|user|>\n{user_msg}</s>\n<|assistant|\n{assistant_msg}</s>\n"
+    conversation += f"<|user|>\n{message}</s>\n<|assistant|"
     
-    return f"Sentiment: {label}\nConfidence: {score:.2%}"
-
-def generate_text(prompt):
-    """Generate text based on the input prompt."""
-    if not prompt.strip():
-        return "Please enter a prompt to generate text."
+    try:
+        # Generate response
+        response = generator(
+            conversation,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=generator.tokenizer.eos_token_id
+        )
+        
+        # Extract only the new response
+        full_text = response[0]['generated_text']
+        # Get the part after the last assistant tag
+        if "<|assistant| " in full_text:
+            assistant_response = full_text.split("<|assistant|")[-1].strip()
+        else:
+            assistant_response = full_text[len(conversation):].strip()
+        
+        # Clean up any remaining tags
+        assistant_response = assistant_response.replace("</s>", "").strip()
+        
+        if not assistant_response:
+            assistant_response = "I'm thinking... could you ask that again?"
+            
+    except Exception as e:
+        assistant_response = f"Sorry, I encountered an error: {str(e)}"
     
-    result = generator(prompt, max_length=100, num_return_sequences=1)
-    return result[0]['generated_text']
+    history.append((message, assistant_response))
+    return history
 
-def chatbot(message, history):
-    """Simple chatbot function."""
-    if not message.strip():
-        return "Please say something!"
+# Create the Gradio interface
+with gr.Blocks(
+    title="AI Chat",
+    theme=gr.themes.Soft(
+        primary_hue="purple",
+        secondary_hue="blue",
+    ),
+    css="""
+    .gradio-container {
+        max-width: 800px !important;
+        margin: auto !important;
+    }
+    """
+) as demo:
+    gr.Markdown("# 🤖 AI Chat Assistant")
+    gr.Markdown(f"Powered by **{MODEL_NAME}**")
     
-    # Simple response logic
-    message_lower = message.lower()
+    chatbot = gr.Chatbot(
+        label="Chat",
+        height=500,
+        show_copy_button=True,
+        bubble_full_width=False,
+        avatar_images=(None, "🤖")
+    )
     
-    if "hello" in message_lower or "hi" in message_lower:
-        return "Hello! How can I help you today?"
-    elif "how are you" in message_lower:
-        return "I'm doing great, thanks for asking! I'm an AI assistant ready to help."
-    elif "bye" in message_lower:
-        return "Goodbye! Have a great day!"
-    elif "sentiment" in message_lower:
-        return "I can analyze sentiment! Go to the Sentiment Analysis tab and enter some text."
-    else:
-        # Use the generator for other inputs
-        response = generator(message, max_length=50, num_return_sequences=1)
-        return response[0]['generated_text']
-
-# Create the Gradio interface with tabs
-with gr.Blocks(title="Basic AI Assistant") as demo:
-    gr.Markdown("# 🤖 Basic AI Assistant")
-    gr.Markdown("A simple AI-powered assistant with multiple capabilities!")
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Message",
+            placeholder="Type your message here...",
+            scale=9,
+            container=False
+        )
+        submit_btn = gr.Button("Send", scale=1, variant="primary")
     
-    with gr.Tabs():
-        with gr.TabItem("💬 Chat"):
-            gr.Markdown("Chat with the AI assistant!")
-            chat_interface = gr.ChatInterface(chatbot, type="messages")
-        
-        with gr.TabItem("😊 Sentiment Analysis"):
-            gr.Markdown("Analyze the sentiment of your text (positive or negative).")
-            with gr.Row():
-                sentiment_input = gr.Textbox(
-                    label="Enter text to analyze",
-                    placeholder="Type something like 'I love this product!'",
-                    lines=3
-                )
-                sentiment_output = gr.Textbox(label="Result", lines=2)
-            sentiment_btn = gr.Button("Analyze Sentiment", variant="primary")
-            sentiment_btn.click(analyze_sentiment, inputs=sentiment_input, outputs=sentiment_output)
-        
-        with gr.TabItem("✍️ Text Generation"):
-            gr.Markdown("Generate text based on a prompt!")
-            with gr.Row():
-                gen_input = gr.Textbox(
-                    label="Enter your prompt",
-                    placeholder="Once upon a time...",
-                    lines=3
-                )
-                gen_output = gr.Textbox(label="Generated Text", lines=5)
-            gen_btn = gr.Button("Generate Text", variant="primary")
-            gen_btn.click(generate_text, inputs=gen_input, outputs=gen_output)
+    with gr.Row():
+        clear_btn = gr.Button("Clear Chat", variant="secondary")
     
-    gr.Markdown("---")
-    gr.Markdown("Made with ❤️ using Gradio and Hugging Face Transformers")
+    # Event handlers
+    msg.submit(chat_with_ai, [msg, chatbot], [chatbot]).then(
+        lambda: "", None, [msg]
+    )
+    submit_btn.click(chat_with_ai, [msg, chatbot], [chatbot]).then(
+        lambda: "", None, [msg]
+    )
+    clear_btn.click(lambda: [], None, [chatbot])
 
 if __name__ == "__main__":
     demo.launch()