Upload app.py

app.py

@@ -8,47 +8,25 @@
 
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import Conversation, pipeline
 
-# Load the pre-trained model and tokenizer
-model_name = 'microsoft/DialoGPT-medium'
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Set the device to GPU if available, otherwise use CPU
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = model.to(device)
-
-# Define a function to generate a response given a list of user inputs
-def generate_response(user_inputs):
-    # Tokenize the user inputs
-    input_ids = tokenizer.encode(user_inputs, return_tensors='pt').to(device)
-
-    # Generate a response
-    with torch.no_grad():
-        output = model.generate(input_ids, max_length=100, num_return_sequences=1)
-
-    # Decode the generated output
-    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return response
+# Load the conversational pipeline
+model_name = "facebook/blenderbot-400M-distill"
+chatbot = pipeline("conversational", model=model_name)
 
 # Define the chatbot function
-def chatbot(input_text):
-    # Append user input to the chat history
-    history = [input_text]
-
-    # Generate a response
-    response = generate_response(history)
+def generate_response(input_text):
+    conversation = Conversation()
+    conversation.add_user_input(input_text)
 
-    # Append the user input and generated response to the chat history
-    history.append(response)
+    response = chatbot(conversation)
 
-    # Return the response
     return response
+    # .choices[0]['message']['content']
 
 # Set up the Gradio interface
 iface = gr.Interface(
-    fn=chatbot,
+    fn=generate_response,
     inputs=gr.inputs.Textbox(placeholder="Enter your message..."),
     outputs="text",
     title="Conversational Chatbot",