import gradio as gr
from transformers import pipeline

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

# Function to load the selected model and classify text
def classify_text(model_name, text):
    classifier = pipeline("text-classification", model=models[model_name], top_k=None)
    predictions = classifier(text)

    # Map the numerical labels to human-readable labels
    label_mapping = {"0": "Male", "1": "Female"}

    # Construct the output dictionary with human-readable labels
    output_predictions = {}
    for pred in predictions[0]:
        # Ensure the label is treated as a string for dictionary lookup
        numerical_label_str = str(pred["label"])
        human_readable_label = label_mapping.get(numerical_label_str, numerical_label_str) # Use fallback if label not in mapping
        output_predictions[human_readable_label] = pred["score"]

    return output_predictions

# Create the Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Dropdown(
            list(models.keys()),
            label="Select Model",
            value="ModernBERT Large (gender)"
        ),
        gr.Textbox(
            lines=2,
            placeholder="Enter text to analyze emotions...",
            value="I am thrilled to be a part of this amazing journey!"
        )
    ],
    outputs=gr.Label(num_top_classes=5),
    title="ModernBERT gender Classifier",
    description="Select a model and enter a sentence to see its associated gender and confidence scores.",
)

# Launch the app
interface.launch()