from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
import gradio as gr

model_name = "gravitee-io/bert-small-toxicity"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

model.eval()
labels = model.config.id2label

def detecter_toxicite(texte):
    """
    Détecte les types de toxicité dans un texte donné en utilisant le modèle BERT.
    Retourne un dictionnaire avec chaque label de toxicité et son score de probabilité.
    """
    # 3. Préparer l'entrée (Tokenization)
    # Le modèle BERT attend des identifiants de tokens et un masque d'attention
    inputs = tokenizer(texte, return_tensors="pt", truncation=True, padding=True, max_length=512)

    # 4. Exécuter l'inférence (sur CPU par défaut)
    with torch.no_grad(): # Désactiver le calcul des gradients pour l'inférence
        outputs = model(**inputs)

    # 5. Post-traitement des sorties
    # Les sorties sont des logits (scores bruts). Nous devons les convertir en probabilités.
    logits = outputs.logits
    probabilities = F.softmax(logits, dim=1)

    resultats = {}
    output = "Scores de toxicité :\n"
    for i, label_id in enumerate(model.config.id2label):
        label_name = model.config.id2label[label_id]
        score = probabilities[0][i].item() # .item() pour obtenir la valeur Python standard

        output += f"  - {label_name}: {score:.4f}\n"
        resultats[label_name.lower()] = score

    # Exemple de décision basée sur un seuil
    seuil_toxicite = 0.2 # Définissez votre propre seuil
    est_toxique = resultats["toxic"] > resultats["not-toxic"]
    if est_toxique:
        output += f"Verdict : TOXIQUE (score > {seuil_toxicite})\n"
    else:
        output += f"Verdict : NON TOXIQUE (aucun score > {seuil_toxicite})\n"
    
    return output

with gr.Blocks() as ui:
    gr.Markdown("# Test modération")

    with gr.Tab("Modération de texte"):
        text_input = gr.Textbox(label="Prompt", placeholder="Je suis le text a modérer", lines=4, scale=2)
        text_output = gr.Textbox(label="Résultat Modération", lines=4, scale=2)
        text_btn = gr.Button("Générer")
        text_btn.click(detecter_toxicite, inputs=text_input, outputs=text_output)

ui.launch()