Changed Structure

app.py

@@ -2,71 +2,35 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# Load DialoGPT model and tokenizer
-model_name = "microsoft/DialoGPT-large"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load the DialoGPT model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
+model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
 
-# Respond function for Gradio interface
-def respond(
-    message,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    # Initialize history if it is None
-    if history is None:
-        history = []
+# Respond function
+def respond(message, chat_history=None):
+    if chat_history is None:
+        chat_history = []
 
-    try:
-        # Format the chat history for the DialoGPT model using the 'messages' format
-        conversation = [{"role": "system", "content": system_message}]
-        for user_msg, bot_msg in history:
-            if user_msg:
-                conversation.append({"role": "user", "content": user_msg})
-            if bot_msg:
-                conversation.append({"role": "assistant", "content": bot_msg})
-        conversation.append({"role": "user", "content": message})
+    # Encode the user input and append to the chat history
+    new_user_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
+    bot_input_ids = torch.cat([torch.tensor(chat_history), new_user_input_ids], dim=-1) if chat_history else new_user_input_ids
 
-        # Tokenize input and generate response
-        inputs = tokenizer.encode(" ".join([msg["content"] for msg in conversation]), return_tensors="pt")
-        outputs = model.generate(
-            inputs,
-            max_length=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-        response = tokenizer.decode(outputs[:, inputs.shape[-1] :][0], skip_special_tokens=True)
+    # Generate the bot's response
+    chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    bot_message = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
 
-        # Append the new user-bot interaction to the history
-        history.append((message, response))
+    # Update chat history
+    chat_history = chat_history_ids.tolist()
 
-        return response, history
-    except Exception as e:
-        # Return the actual error message if something goes wrong
-        return f"Error occurred: {str(e)}", history
+    return bot_message, chat_history
 
-# Gradio Chat Interface
+# Gradio Interface
 demo = gr.Interface(
     fn=respond,
-    inputs=[
-        gr.Textbox(label="Message"),
-        gr.State(),  # For history
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-    outputs=["text", gr.State()],  # Return the response and updated history
+    inputs=["text", gr.State()],
+    outputs=["text", gr.State()],
+    title="DialoGPT Chatbot",
+    description="A chatbot powered by Microsoft's DialoGPT.",
 )
 
 if __name__ == "__main__":