import gradio as gr
from transformers import pipeline

# Load smaller model to reduce latency
qa_pipeline = pipeline("text-generation", model="distilgpt2", device=0)  # Use GPU if available

def real_estate_chatbot(query):
    # Prompt for the chatbot, ensuring the query is real-estate-related
    prompt = f"User: {query}\nAssistant (Real-Estate):"
    
    # Generate a response from the model, limiting the length for faster results
    response = qa_pipeline(prompt, max_length=100, num_return_sequences=1)[0]['generated_text']
    
    # Extract and return the response after the "Assistant (Real-Estate):" part
    return response.split("Assistant (Real-Estate):")[1].strip()

# Create the Gradio interface
input_component = gr.Textbox(placeholder="Ask about buying, renting, or selling properties...")

output_component = gr.Textbox()

# Launch the Gradio interface with caching enabled to speed up repeated queries
interface = gr.Interface(fn=real_estate_chatbot, 
                         inputs=input_component, 
                         outputs=output_component, 
                         title="Real-Estate Chatbot", 
                         description="Ask questions related to buying, selling, or renting properties.",
                         cache_examples=True)

interface.launch()