#data set from Gender by Name [Dataset]. (2020). UCI Machine Learning Repository. https://doi.org/10.24432/C55G7X.
import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import os
model_paths = {
    "bert": "./weights/bert_gender",
    "distilbert": "./weights/distilbert_gender"
}
labels = ["Male", "Female"]

models = {}
tokenizers = {}

for name, path in model_paths.items():
    if os.path.exists(path):
        print(f"Loading {name} from {path}...")
        tokenizers[name] = AutoTokenizer.from_pretrained(path)
        # Use num_labels=2 to match your trained weights
        models[name] = AutoModelForSequenceClassification.from_pretrained(path, num_labels=2)
        models[name].eval()
    else:
        print(f"Warning: Path {path} not found. Did you finish training?")

def predict(text):
    all_probs = []
    #dont keep track of gradients, just loading for demo purposes
    with torch.no_grad():
        for name in models.keys():
            #tokenize and predict
            inputs = tokenizers[name](text, return_tensors="pt", truncation=True, padding=True)
            outputs = models[name](**inputs)
            
            #converts logits to readable probabilities
            probs = F.softmax(outputs.logits, dim=-1).numpy()[0]
            all_probs.append(probs)
    
    #ensemble based off probability
    ensemble_probs = np.mean(all_probs, axis=0)
    
    # Return a dictionary of {label: probability} for label components
    return {labels[i]: float(ensemble_probs[i]) for i in range(len(labels))}

#gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=2, placeholder="Enter a name (e.g., Alex, Jordan)...", label="Input Name"),
    outputs=gr.Label(num_top_classes=2, label="Gender Prediction"),
    title="UCI Gender-by-Name Ensemble",
    description="A voting ensemble of BERT, RoBERTa, and DistilBERT trained on the UCI Gender by Name dataset."
)

if __name__ == "__main__":
    demo.launch()