import gradio as gr
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import torch
from torch.nn import functional as F

# Load the trained intent classification model
model_name = "sonisatish119/PhysioMindAI-intent-classification-bert" 
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name)

# Define intent label mapping
intent_labels = {'reschedule_appointment': 0, 'appointment_location_details': 1, 'cancel_appointment': 2, 'emergency_booking': 3, 'check_appointment_status': 4, 'modify_appointment_details': 5, 'available_slots_inquiry': 6, 'appointment_requirements': 7, 'book_appointment': 8, 'appointment_reminder': 9}

# Function to predict intent
def predict_intent(text):
    model.eval()  # Set to evaluation mode
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=64)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class = torch.argmax(F.softmax(logits, dim=1)).item()
    intent_name = [k for k, v in intent_labels.items() if v == predicted_class][0]
    return intent_name

# Gradio Interface
demo = gr.Interface(
    fn=predict_intent,
    inputs=gr.Textbox(placeholder="Type your query here...", lines=2),
    outputs=gr.Textbox(label="Predicted Intent"),
    title="📝 Intent Classification Model",
    description="This model classifies user queries into predefined appointment-related intents."
)

# Launch the app
if __name__ == "__main__":
    demo.launch()