import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch

# Define model names
models = {
    "gte base (gender v3.1)": "breadlicker45/gte-gender-v3.1-test",
    "ModernBERT Large (gender v3)": "breadlicker45/modernbert-gender-v3-test",
    "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
label_map = {
    "LABEL_0": "Male (0)",
    "0": "Male (0)",
    "LABEL_1": "Female (1)",
    "1": "Female (1)"
}

# A cache to store loaded models/pipelines to speed up subsequent requests
model_cache = {}

# Determine the device to run on (GPU if available, otherwise CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")


# The main classification function, now handles both model types
def classify_text(model_name, text):
    try:
        processed_results = {}
        model_id = models[model_name]

        # --- SPECIAL HANDLING FOR THE GTE MODEL ---
        if "gte-gender" in model_id:
            # Check if model/tokenizer is already in our cache
            if model_id not in model_cache:
                print(f"Loading GTE model and tokenizer manually: {model_id}...")
                tokenizer = AutoTokenizer.from_pretrained(model_id)
                model = AutoModelForSequenceClassification.from_pretrained(model_id, trust_remote_code=True).to(device)
                model_cache[model_id] = (model, tokenizer) # Cache both
            
            model, tokenizer = model_cache[model_id]

            # Tokenize the input text and move to the correct device
            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)

            # Get model predictions
            with torch.no_grad():
                logits = model(**inputs).logits
            
            # Convert logits to probabilities using softmax
            probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
            
            # Format results to match the pipeline's output style
            processed_results[label_map["LABEL_0"]] = probabilities[0].item()
            processed_results[label_map["LABEL_1"]] = probabilities[1].item()

        # --- STANDARD HANDLING FOR PIPELINE-COMPATIBLE MODELS ---
        else:
            # Check if the pipeline is already in our cache
            if model_id not in model_cache:
                print(f"Loading pipeline for model: {model_id}...")
                # Load and cache the pipeline
                model_cache[model_id] = pipeline(
                    "text-classification", 
                    model=model_id, 
                    top_k=None,
                    device=device # Use the determined device
                )
            
            classifier = model_cache[model_id]
            predictions = classifier(text)

            # Process predictions to use friendly labels
            if predictions and isinstance(predictions, list) and predictions[0]:
                for pred in predictions[0]:
                    raw_label = pred["label"]
                    score = pred["score"]
                    friendly_label = label_map.get(raw_label, raw_label)
                    processed_results[friendly_label] = score
        
        return processed_results

    except Exception as e:
        print(f"Error: {e}")
        # Return an error message suitable for gr.Label or gr.JSON
        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="gte base (gender v3.1)" # Default model
        ),
        gr.Textbox(
            lines=2,
            placeholder="Enter text to classify for perceived gender...",
            value="This is an example sentence."
        )
    ],
    # Since we now consistently return a dictionary of {label: score},
    # we can go back to using the nicer-looking gr.Label component!
    outputs=gr.Label(num_top_classes=2, label="Classification Results"),
    title="ModernBERT & GTE 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.",
    allow_flagging="never",
)

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