Update app.py

app.py

@@ -1,23 +1,41 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import gradio as gr
-import os  # ✅ Import os to read environment variables
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load the model and tokenizer
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, torch_dtype=torch.float16, device_map="auto"
+)
 
-# ✅ Get the API token securely from Hugging Face Secrets
-HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
+# Chat function
+def chat(user_input, history):
+    messages = [{"role": "user", "content": user_input}]
+    encoded = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
 
-# Replace this with your model
-MODEL_NAME = "mistralai/Mistral-7B-Instruct"
+    # Generate response
+    with torch.no_grad():
+        output_ids = model.generate(encoded, max_new_tokens=256, do_sample=True)
+    
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    
+    # Append to chat history
+    history.append((user_input, response))
+    return history
 
-# Load tokenizer & model with authentication
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_AUTH_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, token=HF_AUTH_TOKEN)
+# Create Gradio Chatbot UI
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Mistral-7B Chatbot")
+    
+    chatbot = gr.Chatbot()
+    user_input = gr.Textbox(label="Type your message here...")
+    
+    def respond(message, chat_history):
+        return chat(message, chat_history)
 
-# Define Gradio chat function
-def chat_with_ai(message):
-    inputs = tokenizer(message, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=500)
-    return tokenizer.decode(outputs[0])
+    user_input.submit(respond, [user_input, chatbot], chatbot)
 
-# Create Gradio Interface
-iface = gr.Interface(fn=chat_with_ai, inputs="text", outputs="text")
-iface.launch()
+# Launch for Hugging Face Spaces
+demo.launch()