import gradio as gr
from transformers import pipeline

# Define model names
models = {
    "ModernBERT Large (gender v2)": "breadlicker45/modernbert-gender-v2",
    "ModernBERT Base (gender)": "breadlicker45/ModernBERT-base-gender",
    "ModernBERT Large (gender)": "breadlicker45/ModernBERT-large-gender"
}

# 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": "Male (0)",
    "0": "Male (0)",
    "LABEL_1": "Female (1)",
    "1": "Female (1)"
}

# Function to load the selected model and classify text
def classify_text(model_name, text):
    try:
        classifier = pipeline("text-classification", model=models[model_name], top_k=None)
        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="ModernBERT Large (gender)" # Default model
        ),
        gr.Textbox(
            lines=2,
            placeholder="Enter text to classify for perceived gender...", # 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 Gender Classifier",
    description="Select a model and enter a sentence to see the perceived gender classification (Male=0, Female=1) and confidence scores. Note: Text-based gender classification can be unreliable and reflect societal biases.", # Updated description
)

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