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

# Class ID to category name mapping
categories = [
    "FAQ",
    "Escalation",
    "Fallback"
]

# Load tokenizer from the uploaded HF repo
tokenizer = AutoTokenizer.from_pretrained("PVK-Varma/DistilBERT_Banking")

# Load base model and LoRA weights from HF repo
config = AutoConfig.from_pretrained("distilbert-base-uncased", num_labels=len(categories))
base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", config=config)

model = PeftModel.from_pretrained(base_model, "PVK-Varma/DistilBERT_Banking")

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class_id = predictions.argmax().item()
        confidence = predictions.max().item()
        category_name = categories[predicted_class_id]
    return category_name

# Create Gradio interface
iface = gr.Interface(fn=predict, inputs="text", outputs="text", title="DistilBERT Text Classifier")
iface.launch()