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()