Update app.py

app.py

@@ -1,30 +1,56 @@
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
 import os
-from transformers import AutoTokenizer
 
+# 1. Ensure tokenizer exists (Mistral-7B compatible)
 if not os.path.exists("BuddAi/tokenizer.json"):
+    os.makedirs("BuddAi", exist_ok=True)
     tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
     tokenizer.save_pretrained("BuddAi")
-    
-# Load your model (replace with your model ID)
-model_id = "CaptMetal/BuddAi"
+    print("Saved tokenizer files to BuddAi/")
+
+# 2. Load model and tokenizer
+model_id = "BuddAi"  # Local path (or your HF repo "CaptMetal/BuddAi" if uploaded)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id)
-chatbot = pipeline("conversational", model=model, tokenizer=tokenizer)
 
+# 3. Create pipeline with proper chat template
+tokenizer.chat_template = "{% for message in messages %}{{message['content']}}{% if not loop.last %}{{' '}}{% endif %}{% endfor %}"
+chatbot = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device_map="auto"  # Uses GPU if available
+)
 
-    
-# Chat function with memory
+# 4. Improved chat function
 def respond(message, history):
-    conversation = []
+    # Format conversation history
+    formatted_history = ""
     for user_msg, bot_msg in history:
-        conversation.append({"role": "user", "content": user_msg})
-        conversation.append({"role": "assistant", "content": bot_msg})
-    conversation.append({"role": "user", "content": message})
+        formatted_history += f"<|user|>{user_msg}</s><|assistant|>{bot_msg}</s>"
+    
+    # Current message
+    prompt = f"{formatted_history}<|user|>{message}</s><|assistant|>"
+    
+    # Generate response
+    outputs = chatbot(
+        prompt,
+        max_new_tokens=256,
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
     
-    response = chatbot(conversation)
-    return response[-1]["generated_text"]
+    # Extract and clean response
+    full_text = outputs[0]["generated_text"]
+    response = full_text[len(prompt):].split("</s>")[0].strip()
+    return response
 
-# Launch interface
-gr.ChatInterface(respond).launch()
+# 5. Launch Gradio interface
+gr.ChatInterface(
+    respond,
+    title="BuddAI - Mistral-7B Chatbot",
+    description="A conversational AI friend powered by Mistral-7B",
+    examples=["How are you today?", "Tell me a joke!"]
+).launch(server_port=7860, share=True)