Update app.py

app.py

@@ -1,41 +1,23 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
-import torch
+from transformers import pipeline
 
-# Load the model and tokenizer
-model_name = "microsoft/DialoGPT-small"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load the pipeline for text generation
+generator = pipeline("text-generation", model="microsoft/DialoGPT-small")
 
 # 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([history_tensor, new_user_input_ids], dim=-1)
-
-    # Generate a response
-    chat_history_ids = model.generate(
-        bot_input_ids, 
-        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, keeping only the new tokens
-    response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    # Since the pipeline handles everything, we just need to format our input
+    conversation = [{"role": "user", "content": user_input}] if history is None else history + [{"role": "user", "content": user_input}]
     
-    # Update history with new input and response
-    new_history = chat_history_ids.tolist()[0]  # Convert tensor to list for Gradio state
+    # Generate response using the pipeline, which manages all pre/post-processing
+    response = generator(conversation, return_full_text=False, max_length=1000)
 
-    return response, new_history
+    # Extract the last assistant response
+    assistant_response = response[0]['generated_text']
+    
+    # Append this response to history
+    new_history = conversation + [{"role": "assistant", "content": assistant_response}]
+    
+    return assistant_response, new_history
 
 # Gradio interface
 iface = gr.Interface(
@@ -50,7 +32,7 @@ iface = gr.Interface(
     ],
     title="DialoGPT Chat",
     description="Chat with DialoGPT-small model. Your conversation history is maintained.",
-    allow_flagging="never"  # Disabling flagging since this is a chat model
+    allow_flagging="never"
 )
 
 iface.launch()