Update app.py

app.py

@@ -2,7 +2,6 @@ from transformers import GPT2LMHeadModel, AutoTokenizer, pipeline
 import gradio as gr
 import torch
 
-# Load model from Hugging Face Hub
 model = GPT2LMHeadModel.from_pretrained("GBhaveshKumar/ConvoAI")
 tokenizer = AutoTokenizer.from_pretrained("GBhaveshKumar/ConvoAI")
 
@@ -13,9 +12,12 @@ generator = pipeline(
     device=0 if torch.cuda.is_available() else -1
 )
 
-def chat_with_bot(history, user_input):
-    context = "".join([f"A: {q}\nB: {a}\n" for q, a in history])
-    context += f"A: {user_input}\nB:"
+def respond(message, chat_history):
+    context = ""
+    for user, bot in chat_history:
+        context += f"A: {user}\nB: {bot}\n"
+    context += f"A: {message}\nB:"
+
     output = generator(
         context,
         max_length=len(tokenizer.encode(context)) + 50,
@@ -25,8 +27,13 @@ def chat_with_bot(history, user_input):
         top_p=0.95,
         temperature=0.8
     )[0]['generated_text']
+
     reply = output[len(context):].split("\n")[0].strip()
-    history.append((user_input, reply))
-    return history, ""
+    chat_history.append((message, reply))
+    return "", chat_history
 
-gr.ChatInterface(fn=chat_with_bot, title="💬 DailyDialog Chatbot").launch()
+gr.ChatInterface(
+    fn=respond,
+    title="💬 DailyDialog Chatbot",
+    description="A fine-tuned GPT-2 chatbot trained on the DailyDialog dataset.",
+).launch()