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 and their scores
    top_labels = result["labels"][:2]
    scores = result["scores"][:2]
    
    # Prepare the outputs to display both labels and their corresponding meters
    outputs = []
    for label, score in zip(top_labels, scores):
        label_with_score = f"{label}: {score:.2f}"
        outputs.append(gr.Label(label_with_score))
        outputs.append(gr.Meter(value=score))
    
    return outputs

# Create Gradio interface
feedback_textbox = gr.Textbox(label="Enter your feedback:")
feedback_output = [gr.Label(), gr.Meter(), gr.Label(), gr.Meter()]  # Output placeholders for two labels and meters

iface = 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.",
    layout="vertical"
)

iface.launch()