Update app.py

app.py

@@ -7,48 +7,43 @@ model_name = "microsoft/DialoGPT-medium"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-chat_history_ids = None
-
-# Persona templates
+# Persona options (used for prompt formatting)
 persona_prompts = {
-    "Elon Musk": "You are Elon Musk, respond as a visionary tech entrepreneur with bold ideas and futuristic thinking.",
-    "Jensen Huang": "You are Jensen Huang, respond with deep technical insight, leadership wisdom, and references to GPUs and AI innovation.",
-    "Jeff Bezos": "You are Jeff Bezos, respond as a strategic business titan with a calm, calculated, and confident tone."
+    "Elon Musk": "Elon Musk",
+    "Jensen Huang": "Jensen Huang",
+    "Jeff Bezos": "Jeff Bezos"
 }
 
-# NOTE: Function parameters swapped to match Gradio input order
 def chatbot(persona, input_text):
-    global chat_history_ids
-    # Inject persona
-    persona_instruction = persona_prompts.get(persona, "")
-    full_input = f"{persona_instruction}\nUser: {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
-    new_input_ids = tokenizer.encode(full_input + tokenizer.eos_token, return_tensors='pt')
-    # Combine with history
-    bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1) if chat_history_ids is not None else new_input_ids
-    # Create attention mask
-    attention_mask = torch.ones(bot_input_ids.shape, dtype=torch.long)
+    input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors='pt')
+
     # Generate response
-    chat_history_ids = model.generate(
-        bot_input_ids,
-        attention_mask=attention_mask,
+    output_ids = model.generate(
+        input_ids,
         max_length=1000,
         pad_token_id=tokenizer.eos_token_id
     )
-    response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return response
 
-# Gradio interface
+    # 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
 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...")
+        gr.Textbox(lines=2, placeholder="Say something...", label="input_text")
     ],
     outputs="text",
     title="Persona Bot (DialoGPT)",
     description="Chat with the voice of Elon Musk, Jensen Huang, or Jeff Bezos. Powered by Hugging Face + Transformers."
 )
 
-# Launch with public sharing
+# Launch with share enabled
 iface.launch(share=True)