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 = ["Procedures", "Maintenance", "Operations", "Health", "Safety"]

    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]
    
    # Check if the accuracy of the top label is less than 30%
    if scores[0] < 0.3:
        return "Please provide another relevant feedback."
    
    # Generate HTML content for displaying the scores as meters/progress bars
    html_content = ""
    for i in range(len(top_labels)):
        score_percentage = scores[i] * 100  # Convert score to percentage
        html_content += f"<div><b>{top_labels[i]}:</b> {scores[i]:.2f} <div style='background-color: #e0e0e0; border-radius: 10px;'><div style='height: 24px; width: {score_percentage}%; background-color: #76b900; border-radius: 10px;'></div></div></div>"
    
    return html_content

# Create Gradio interface
feedback_textbox = gr.Textbox(label="Enter your feedback:")
feedback_output = gr.HTML(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."
).launch()