import torch
from transformers import BertTokenizer, BertForSequenceClassification
import gradio as gr

# Load model and tokenizer from local directory (same folder as app.py)
model = BertForSequenceClassification.from_pretrained("bert")
tokenizer = BertTokenizer.from_pretrained(".")

# Ensure model runs on GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# ID to label mapping
id2label = {0: "Select", 1: "Insert", 2: "Delete", 3: "Update", 4: "Analyse"}

def classify_query(text):
    if not text.strip():
        return "Please enter a valid query."
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        prediction = torch.argmax(outputs.logits, dim=1).item()
        label = id2label.get(prediction, "Unknown")
    return f"🧠 Predicted Query Type: **{label}**"

# Gradio UI
demo = gr.Interface(
    fn=classify_query,
    inputs=gr.Textbox(label="Enter your expense query", lines=2, placeholder="e.g., Show me all expenses from January."),
    outputs=gr.Markdown(label="Query Type"),
    title="💰 Expense Query Type Classifier",
    description="This model classifies your natural language query into one of 5 SQL operation types: Select, Insert, Delete, Update, or Analyse.",
    examples=[
        ["Add an expense of 500 in groceries at Amazon"],
        ["Remove last transaction from Starbucks"],
        ["Update amount of food expense to 850"],
        ["Kitna kharcha hua electronics par?"],
        ["Give me analytics of travel spending"]
    ],
    theme="soft",
)

if __name__ == "__main__":
    demo.launch()