Update app.py

app.py

@@ -1,49 +1,53 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Load model and tokenizer
-model_name = "microsoft/DialoGPT-medium"
+model_name = "openchat/openchat-3.5-1210"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
-# Persona options (used for prompt formatting)
+# Persona prompts
 persona_prompts = {
-    "Elon Musk": "Elon Musk",
-    "Jensen Huang": "Jensen Huang",
-    "Jeff Bezos": "Jeff Bezos"
+    "Elon Musk": "You are Elon Musk, a visionary tech entrepreneur with bold ideas and futuristic thinking.",
+    "Jensen Huang": "You are Jensen Huang, an AI hardware leader with deep technical insights and a passion for GPUs.",
+    "Jeff Bezos": "You are Jeff Bezos, a calculated and confident business titan with strategic thinking."
 }
 
-def chatbot(persona, input_text):
-    # Format as conversation between User and chosen persona
-    speaker = persona_prompts.get(persona, "Person")
-    prompt = f"User: {input_text}\n{speaker}:"
-
-    # Encode input
-    input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors='pt')
-
-    # Generate response
-    output_ids = model.generate(
-        input_ids,
-        max_length=1000,
-        pad_token_id=tokenizer.eos_token_id
-    )
+def format_openchat_prompt(system_prompt, user_input):
+    return f"<|system|>\n{system_prompt}\n<|end|>\n<|user|>\n{user_input}\n<|end|>\n<|assistant|>\n"
 
-    # Decode and return only the new response
-    output = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return f"{speaker}: {output.strip()}"
-
-# Gradio Interface
+def chatbot(persona, input_text):
+    system_prompt = persona_prompts.get(persona, "")
+    prompt = format_openchat_prompt(system_prompt, input_text)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id
+        )
+
+    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    assistant_response = full_output.split("<|assistant|>")[-1].strip()
+    return assistant_response
+
+# Gradio UI
 iface = gr.Interface(
     fn=chatbot,
     inputs=[
-        gr.Dropdown(choices=["Elon Musk", "Jensen Huang", "Jeff Bezos"], label="Choose Persona"),
-        gr.Textbox(lines=2, placeholder="Say something...", label="input_text")
+        gr.Dropdown(choices=list(persona_prompts.keys()), label="Choose Persona"),
+        gr.Textbox(lines=2, placeholder="Ask something...")
     ],
     outputs="text",
-    title="Persona Bot (DialoGPT)",
-    description="Chat with the voice of Elon Musk, Jensen Huang, or Jeff Bezos. Powered by Hugging Face + Transformers."
+    title="Persona Bot (OpenChat)",
+    description="Chat with the voice of Elon Musk, Jensen Huang, or Jeff Bezos. Powered by OpenChat 3.5 + Transformers.",
 )
 
-# Launch with share enabled
 iface.launch(share=True)