import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("diabolic6045/itineraries_Generator")
model = AutoModelForCausalLM.from_pretrained("diabolic6045/itineraries_Generator")

# Function to generate itinerary
def generate_itinerary(destination, num_days):
    prompt = f"{destination} ({num_days} days): itinerary:"
    inputs = tokenizer(prompt, return_tensors="pt")
    output = model.generate(**inputs, max_length=50)
    itinerary = tokenizer.decode(output[0], skip_special_tokens=True)
    return itinerary

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Travel Itinerary Generator Demo")
    gr.Markdown("Enter your travel destination and the number of days, and get a personalized itinerary.")
    
    destination = gr.Textbox(label="Destination", placeholder="e.g., Tokyo, Japan")
    num_days = gr.Slider(label="Number of Days", minimum=1, maximum=10, step=1, value=5)
    output = gr.Textbox(label="Generated Itinerary")
    
    generate_button = gr.Button("Generate Itinerary")
    
    generate_button.click(generate_itinerary, inputs=[destination, num_days], outputs=output)

# Launch the Gradio app
demo.launch()