import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
from peft import PeftModel
import gradio as gr
import joblib

# Load base model
base_model_name = "hopjetair/intent_new"

# Load base model and adapter
model = AutoModelForSequenceClassification.from_pretrained(base_model_name, num_labels=53)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_name)

# Load label encoder
label_encoder = joblib.load("label_encoder.pkl")

def predict_intent(text):
    clf = pipeline("text-classification", model=model, tokenizer=tokenizer)
    result = clf(text)[0]
    label_num = int(result["label"].split("_")[-1])
    # Convert back to label
    predicted_label = label_encoder.inverse_transform([label_num])[0]
    return predicted_label

# Gradio UI
demo = gr.Interface(
    fn=predict_intent,
    inputs=gr.Textbox(placeholder="Enter a user query, e.g., 'Book a flight to Boston'"),
    outputs="text",
    title="PEFT Intent Recognition Demo",
    description="Enter a natural language command to classify it into one of 33 intents."
)

demo.launch()