import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

model_name = "rahmanazhar/Travereel-Model-V1"

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    itinerary_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

    def generate_kerala_itinerary(duration, interests, budget, specific_places="Kerala"):
        """Generates a Kerala trip itinerary based on user inputs."""
        prompt = f"Generate a {duration}-day trip plan for {specific_places}, Kerala, focusing on {interests} with a {budget} budget."
        generated_text = itinerary_generator(prompt, max_length=500, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95)[0]['generated_text']
        return generated_text.strip()

    # Define the input and output interfaces for Gradio
    iface = gr.Interface(
        fn=generate_kerala_itinerary,
        inputs=[
            gr.Slider(minimum=1, maximum=10, step=1, label="Duration (Days)"),
            gr.Textbox(label="Your Interests (e.g., beaches, mountains, backwaters, culture)"),
            gr.Radio(choices=["low", "medium", "high"], label="Budget"),
            gr.Textbox(label="Specific Places in Kerala (Optional, leave blank for general Kerala trip)")
        ],
        outputs=gr.Textbox(label="Generated Kerala Trip Plan"),
        title="Kerala Trip Planner AI",
        description="Enter your desired trip duration, interests, and budget to get a personalized Kerala itinerary powered by the Travereel model."
    )

    # Launch the Gradio interface
    iface.launch()

except Exception as e:
    print(f"Error loading model: {e}")
    def error_message():
        return f"Error loading the model: {e}. Please check the logs."
    iface = gr.Interface(fn=error_message, inputs=[], outputs=gr.Textbox(), title="Kerala Trip Planner AI (Error)")
    iface.launch()