Update app.py

app.py

@@ -1,5 +1,6 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
+import torch
 
 # Load the model and tokenizer
 model_name = "microsoft/DialoGPT-small"
@@ -8,30 +9,48 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Function to generate a response
 def dialoGPT_response(user_input, history):
+    # Convert history to tensor if it's not None
+    if history:
+        history_tensor = torch.LongTensor(history)
+    else:
+        history_tensor = torch.LongTensor([])
+
     # Encode the new user input, with the history
     new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
 
     # Append the new user input tokens to the chat history
-    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1) if history else new_user_input_ids
+    bot_input_ids = torch.cat([history_tensor, new_user_input_ids], dim=-1)
 
     # Generate a response
     chat_history_ids = model.generate(
         bot_input_ids, 
-        max_length=1000,
-        pad_token_id=tokenizer.eos_token_id
+        max_length=1000,  # You might want to adjust this based on your needs
+        pad_token_id=tokenizer.eos_token_id,
+        no_repeat_ngram_size=3  # This prevents repeating phrases
     )
 
-    # Decode the response
+    # Decode the response, keeping only the new tokens
     response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return response
+    
+    # Update history with new input and response
+    new_history = chat_history_ids.tolist()[0]  # Convert tensor to list for Gradio state
+
+    return response, new_history
 
 # Gradio interface
 iface = gr.Interface(
     fn=dialoGPT_response,
-    inputs=[gr.Textbox(placeholder="Enter your message..."), "state"],
-    outputs="text",
+    inputs=[
+        gr.Textbox(placeholder="Enter your message..."),
+        "state"
+    ],
+    outputs=[
+        "text",  # The response
+        "state"  # Updated history
+    ],
     title="DialoGPT Chat",
-    description="Chat with DialoGPT-small model."
+    description="Chat with DialoGPT-small model. Your conversation history is maintained.",
+    allow_flagging="never"  # Disabling flagging since this is a chat model
 )
 
 iface.launch()