Update app.py

app.py

@@ -1,43 +1,33 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-# Load model and tokenizer
-model_name = "openchat/openchat-3.5-1210"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
+# Load lightweight model
+model_name = "google/flan-t5-small"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Persona prompts
+# Simple persona instructions
 persona_prompts = {
-    "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."
+    "Elon Musk": "As Elon Musk, visionary tech entrepreneur, answer: ",
+    "Jensen Huang": "As Jensen Huang, AI hardware innovator, answer: ",
+    "Jeff Bezos": "As Jeff Bezos, strategic business leader, answer: "
 }
 
-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"
-
 def chatbot(persona, input_text):
-    system_prompt = persona_prompts.get(persona, "")
-    prompt = format_openchat_prompt(system_prompt, input_text)
+    prompt = persona_prompts.get(persona, "") + input_text
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
-
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
+            max_new_tokens=128,
             do_sample=True,
             temperature=0.7,
-            top_p=0.9,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=pad_token_id
+            top_p=0.9
         )
 
-    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    assistant_response = full_output.split("<|assistant|>")[-1].strip()
-    return assistant_response
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Gradio UI
 iface = gr.Interface(
@@ -47,8 +37,8 @@ iface = gr.Interface(
         gr.Textbox(lines=2, placeholder="Ask something...")
     ],
     outputs="text",
-    title="Persona Bot (OpenChat)",
-    description="Chat with the voice of Elon Musk, Jensen Huang, or Jeff Bezos. Powered by OpenChat 3.5 + Transformers.",
+    title="Persona Bot (Fast Mode)",
+    description="Chat quickly with Elon Musk, Jensen Huang, or Jeff Bezos using a lightweight model.",
 )
 
-iface.launch()
+iface.launch(share=True)