import gradio as gr
from transformers import pipeline, BertTokenizer, BertForSequenceClassification

# Charger le modèle zéro-shot de Hugging Face
zero_shot_classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

# Charger le modèle few-shot à partir du modèle sauvegardé
tokenizer = BertTokenizer.from_pretrained('./animal_offense_model')
model = BertForSequenceClassification.from_pretrained('./animal_offense_model')
few_shot_classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

# Fonction pour classifier le texte avec le modèle zéro-shot
def classify_with_zero_shot(input_text):
    candidate_labels = [
        "non-offensive, it's cute! 😇", 
        "very slightly offensive, but not a big deal! 😅", 
        "slightly offensive, just a little! 🤏", 
        "a bit offensive, ouch! 🤭", 
        "moderately offensive, getting there! 😬", 
        "fairly offensive, watch out! 🚨", 
        "offensive, that's a no-no! 🚫", 
        "very offensive, you really shouldn't say that! 😳", 
        "extremely offensive, seriously? 😡", 
        "totally unacceptable and offensive, you are crazy! 🤯"
    ]

    result = zero_shot_classifier(input_text, candidate_labels)
    labels_scores = dict(zip(result["labels"], result["scores"]))
    return labels_scores

# Fonction pour classifier le texte avec votre propre modèle few-shot
def classify_with_few_shot(input_text):
    result = few_shot_classifier(input_text)
    label = result[0]["label"]
    score = result[0]["score"]
    
    # Ajuster la logique pour garantir l'interprétation correcte des labels
    if label == "LABEL_0":
        return {"non-offensive, it's cute! 😇": score, "offensive": 1 - score}
    elif label == "LABEL_1":
        return {"offensive": score, "non-offensive, it's cute! 😇": 1 - score}
    else:
        return {"unknown": 1.0}  # Pour des fins de débogage si un label inattendu est trouvé

# Fonction principale pour sélectionner le modèle
def classify_text(input_text, model_choice):
    if model_choice == "Zero-Shot Model":
        return classify_with_zero_shot(input_text)
    elif model_choice == "Few-Shot Model":
        return classify_with_few_shot(input_text)
    else:
        return "Please select a valid model."

# Liste de phrases exemples (chaque sous-liste est [texte])
example_phrases = [
    ["Your dog is the cutest ever!"],
    ["I think your cat needs to lose some weight."],
    ["Why would anyone like such an ugly fish?"],
    ["Oh no, saying that about a rabbit is not okay at all!"],
    ["That’s a bit harsh on a parrot."],
    ["You should be more gentle when talking about horses."],
    ["This kitten is just too adorable!"],
    ["Wow, calling a bird annoying is really offensive!"],
    ["That’s a lovely compliment for a hamster!"],
    ["Saying that a dog smells bad is quite rude!"]
]

# Créer une interface Gradio Blocks pour plus de flexibilité
with gr.Blocks() as iface:
    gr.Markdown("# Animal Offense Detector")

    with gr.Column():
        gr.Markdown("## Enter Your Text Below:")
        text_input = gr.Textbox(lines=5, placeholder="Enter your text here...")
        model_choice = gr.Radio(choices=["Zero-Shot Model", "Few-Shot Model"], label="Choose Model")
        label_output = gr.Label(label="Labels and Scores")

        gr.Interface(fn=classify_text, inputs=[text_input, model_choice], outputs=label_output)
        
        gr.Examples(examples=example_phrases, inputs=[text_input])

    # Ajouter de l'espacement et une taille de police plus grande pour la documentation
    gr.Markdown("""
    <div style="margin-top: 40px; font-size: 18px;">
    
    #### Documentation for `Animal Offense Detector`

    This script classifies text to determine the level of offense towards animals using two natural language processing models from Hugging Face. Users can choose between a zero-shot model and a few-shot model to evaluate the input text.

    #### Libraries Used
    - **gradio**: Used to create web-based user interfaces for Python functions.
    - **transformers**: Provides machine learning models for natural language processing tasks.

    #### Features
    1. **Zero-Shot Model Classification**: 
       - Uses `facebook/bart-large-mnli` to classify text based on several predefined labels.
       - This model can understand and classify text without needing specific training for each task.
    
    2. **Few-Shot Model Classification**:
       - Uses a custom-trained BERT model to evaluate text as "non-offensive" or "offensive".
       - This model provides a quick and accurate classification based on the training data.

    3. **Model Selection**:
       - The interface allows the user to choose between the zero-shot and few-shot models to classify the text.

    4. **Example Phrases**:
       - Provides example phrases that the user can select to test the models. Each example is designed to test different levels of potential offense.
    </div>
    """)

if __name__ == "__main__":
    iface.launch(share=True)