watchcochatbot / app.py
nitish3031's picture
Update app.py
d44d00e verified
from transformers import pipeline
import gradio as gr
# Load the DialoGPT-medium model
chatbot = pipeline("text-generation", model="codealtgeek/DiabloGPT-medium-rickmorty")
# Detailed context about your watch company
context = """
You are a helpful assistant for WatchCo, a company that sells high-quality watches. Here is some information about our policies:
- Our return policy allows customers to return watches within 30 days for a full refund, as long as the watch is in its original condition.
- We offer a standard 1-year warranty on all watches, covering manufacturing defects.
- Extended warranties are available for purchase, extending coverage to 2 or 3 years.
- We sell three types of watches that are diver watches, dress watches and field watches.
- Prices of our watches start as low as 50 dollars.
- Our average shipping time is 2-weeks.
- There is also an option to get faster delhivery if needed, that is within 1-week.
Please answer the user's question based on this information.
"""
# Function to generate a response
def generate_response(user_input):
prompt = f"{context}\nUser: {user_input}\nAssistant:"
response = chatbot(prompt, max_length=150, num_return_sequences=1)
generated_text = response[0]["generated_text"]
# Extract just the assistant's answer
assistant_response = generated_text.split("Assistant:")[1].strip()
return assistant_response
# Create the chat interface
iface = gr.Interface(fn=generate_response, inputs="text", outputs="text")
iface.launch()