Update app.py

app.py

@@ -1,25 +1,42 @@
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
-# Load a lightweight conversational model (CPU-friendly)
+# Load a lightweight model (CPU-friendly)
 model_name = "microsoft/DialoGPT-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-chatbot = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)  # device=-1 forces CPU
+# Store chat history per session using Gradio's state
+def predict(message, history):
+    # Format conversation for DialoGPT
+    chat_history = ""
+    for human, ai in history:
+        chat_history += f"User: {human}\nBot: {ai}\n"
+    chat_history += f"User: {message}\nBot:"
 
-def respond(message, history):
-    # Format history for DialoGPT
-    history_str = "\n".join([f"User: {h[0]}\nBot: {h[1]}" for h in history])
-    input_text = f"{history_str}\nUser: {message}\nBot:"
+    # Tokenize and generate
+    inputs = tokenizer.encode(chat_history, return_tensors="pt")
+    outputs = model.generate(
+        inputs,
+        max_length=1000,
+        pad_token_id=tokenizer.eos_token_id,
+        do_sample=True,
+        top_p=0.9,
+        top_k=50
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    response = chatbot(input_text, max_length=1000, pad_token_id=tokenizer.eos_token_id)
-    bot_reply = response[0]['generated_text'].split("Bot:")[-1].strip()
+    # Extract only the bot's latest reply
+    bot_reply = response.split("Bot:")[-1].strip()
     return bot_reply
 
+# Launch the chat interface
 gr.ChatInterface(
-    respond,
-    title="My Chatbot",
+    fn=predict,
+    title="💬 My Chatbot",
     description="A simple CPU-friendly chatbot using DialoGPT-small.",
-    examples=["Hello!", "How are you?", "Tell me a joke."]
+    examples=["Hello!", "What's your name?", "Tell me a fun fact."],
+    retry_btn=None,
+    undo_btn=None,
+    clear_btn="Clear",
 ).launch()