import gradio as gr
from transformers import pipeline

# Load the model
model_name = "knowledgator/comprehend_it-base"
classifier = pipeline("zero-shot-classification", model=model_name, device="cpu")

# Function to classify feedback
def classify_feedback(feedback_text):
    # Classify feedback using the loaded model
    labels = ["Value", "Facilities", "Experience", "Functionality", "Quality"]
    result = classifier(feedback_text, labels, multi_label=True)
    
    # Get the top two labels associated with the feedback
    top_labels = result["labels"][:2]
    scores = result["scores"][:2]
    
    return {top_labels[i]: scores[i] for i in range(len(top_labels))}

# Create Gradio interface
feedback_textbox = gr.Textbox(label="Enter your feedback:")
feedback_output = gr.Textbox(label="Top 2 Labels with Scores:")

gr.Interface(
    fn=classify_feedback,
    inputs=feedback_textbox,
    outputs=feedback_output,
    title="Feedback Classifier",
    description="Enter your feedback and get the top 2 associated labels with scores.",
    capture_session=True
).launch()