Update app.py

app.py

@@ -1,56 +1,28 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import gradio as gr
-import os
 
-# 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")
-    print("Saved tokenizer files to BuddAi/")
-
-# 2. Load model and tokenizer
-model_id = "BuddAi"  # Local path (or your HF repo "CaptMetal/BuddAi" if uploaded)
+# Load OpenHermes
+model_id = "BuddAi"  # Local folder
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+
+# Set chat template (OpenHermes-specific)
+tokenizer.chat_template = """{% for message in messages %}
+{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>'}}
+{% endfor %}"""
 
-# 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
+    temperature=0.7,
+    max_new_tokens=200
 )
 
-# 4. Improved chat function
 def respond(message, history):
-    # Format conversation history
-    formatted_history = ""
-    for user_msg, bot_msg in history:
-        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
-    )
-    
-    # Extract and clean response
-    full_text = outputs[0]["generated_text"]
-    response = full_text[len(prompt):].split("</s>")[0].strip()
-    return response
+    messages = [{"role": "user", "content": message}]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+    outputs = chatbot(prompt)
+    return outputs[0]["generated_text"][len(prompt):].split("<|im_end|>")[0]
 
-# 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)
+gr.ChatInterface(respond).launch()