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

print("Loading Shasa model...")
base_model = AutoModelForCausalLM.from_pretrained(
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    torch_dtype=torch.float32,
)
model = PeftModel.from_pretrained(base_model, "madhan0809/shasa-travel-ai-v0.1")
model = model.merge_and_unload()
model.eval()
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print("Model loaded!")

SYSTEM = "You are NxVoy Shasa, an expert travel AI. Generate detailed trip itineraries with flights, hotels, daily activities, meals, and budget breakdown."

def generate(destination, duration, budget, currency, trip_type, interests):
    prompt = f"<|im_start|>system\n{SYSTEM}<|im_end|>\n<|im_start|>user\nPlan a {int(duration)}-day trip to {destination}. Budget: {int(budget)} {currency}. Trip type: {trip_type}. Interests: {interests}.<|im_end|>\n<|im_start|>assistant\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs, max_new_tokens=1024, temperature=0.4,
            do_sample=True, top_p=0.9, repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
        )
    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True)

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Destination", value="Bali, Indonesia"),
        gr.Number(label="Duration (days)", value=7),
        gr.Number(label="Budget", value=3000),
        gr.Dropdown(["GBP", "USD", "EUR"], label="Currency", value="GBP"),
        gr.Dropdown(["couple_romantic", "family", "solo_adventure", "backpacking"], label="Trip Type", value="couple_romantic"),
        gr.Textbox(label="Interests", value="beaches, temples, food"),
    ],
    outputs=gr.Textbox(label="Shasa Itinerary", lines=20),
    title="Shasa Travel AI by NxVoy",
    description="Shasa is NxVoy's fine-tuned travel AI. Enter trip details to generate a personalized itinerary.",
    examples=[
        ["Bali, Indonesia", 7, 3000, "GBP", "couple_romantic", "beaches, temples, food"],
        ["Tokyo, Japan", 10, 5000, "USD", "solo_adventure", "history, anime, street food"],
        ["Paris, France", 5, 2000, "EUR", "couple_romantic", "art, wine, cuisine"],
    ],
)

demo.launch()