import gradio as gr
from transformers import pipeline

# Define model names
models = {
    "ModernBERT Slop Classifier v1": "underscore2/modernbert_base_slop_classifier",
    "ModernBERT Slop Classifier v2": "underscore2/modernbert_base_slop_classifier_v2",
    "ModernBert Large Slop Classifier v3 (Best accuracy)": "underscore2/modernbert_large_slop_classifier_v3"
}

# Define the mapping for user-friendly labels
# Note: Transformers pipelines often output 'LABEL_0', 'LABEL_1'.
# We handle potential variations like just '0', '1'.
label_map = {
    "LABEL_0": "Human (0)",
    "0": "Human (0)",
    "LABEL_1": "LLM (1)",
    "1": "LLM (1)"
}
large_v3 = pipeline("text-classification", model="underscore2/modernbert_large_slop_classifier_v3", top_k=None)
       
# Function to load the selected model and classify text
def classify_text(model_name, text):
    try:
        if models[model_name] != "underscore2/modernbert_large_slop_classifier_v3":
            classifier = pipeline("text-classification", model=models[model_name], top_k=None)
        else:
            classifier = large_v3
        predictions = classifier(text)

        # Process predictions to use friendly labels
        processed_results = {}
        if predictions and isinstance(predictions, list) and predictions[0]:
             # predictions[0] should be a list of label dicts like [{'label': 'LABEL_1', 'score': 0.9...}, ...]
            for pred in predictions[0]:
                raw_label = pred["label"]
                score = pred["score"]
                # Use the map to get a friendly name, fallback to the raw label if not found
                friendly_label = label_map.get(raw_label, raw_label)
                processed_results[friendly_label] = score
        return processed_results
    except Exception as e:
        # Handle potential errors during model loading or inference
        print(f"Error: {e}")
        # Return an error message suitable for gr.Label
        return {"Error": f"Failed to process: {e}"}


# Create the Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Dropdown(
            list(models.keys()),
            label="Select Model",
            value=list(models.keys())[2] # Default model
        ),
        gr.Textbox(
            lines=2,
            placeholder="Enter text to classify", # Corrected placeholder
            value="This is an example sentence." # Changed example text
        )
    ],
    # The gr.Label component works well for showing classification scores
    outputs=gr.Label(num_top_classes=2), # Show both classes explicitly
    title="ModernBERT Slop Classifier",
    description="Enter a sentence to see the slop and confidence scores", # Updated description
)

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