!pip install -U -q transformers accelerate bitsandbytes kagglehub
## !pip install torch numpy pandas polars

import polars as pl

df_train = pl.read_csv(
    "https://raw.githubusercontent.com/mathisdrn/nlp-representation-learning/refs/heads/main/data/nli_fr_train.tsv",
    separator="\t",
    new_columns=["premise", "hypothesis", "label"],
)
df_val = pl.read_csv(
    "https://raw.githubusercontent.com/mathisdrn/nlp-representation-learning/refs/heads/main/data/nli_fr_test.tsv",
    separator="\t",
    new_columns=["premise", "hypothesis", "label"],
)
label_mapping = {
    "neutral": "Neutre",
    "entailment": "Conséquence",
    "contradiction": "Contradiction",
}
df_train = df_train.with_columns(pl.col("label").replace(label_mapping))
df_val = df_val.with_columns(pl.col("label").replace(label_mapping))
df_train, df_val

zero_shot_prompt = (
    "Tu es un expert en logique. "
    "Classifie la relation entre la prémisse et l'hypothèse en utilisant uniquement : "
    "Conséquence, Contradiction, ou Neutre.\n\n"
    "Classifie l'exemple suivant :\n"
)

# Define the template structure
example_template = pl.format(
    "Prémisse : {}\nHypothèse : {}\nClassification : {}\n---",
    pl.col("premise"),
    pl.col("hypothesis"),
    pl.col("label"),
)

# Extract a random example of each label
examples_str = (
    df_train.group_by("label")
    .tail(1)
    .select(example_template)
    .to_series()
    .str.join("\n")
    .item()
)

few_shot_prompt = (
    "Tu es un expert en logique. "
    + "Classifie la relation entre la prémisse et l'hypothèse en utilisant uniquement : "
    + "Conséquence, Contradiction, ou Neutre.\n\n"
    + "Voici des exemples de classification corrects :\n"
    + examples_str
    + "Réponds maintenant pour l'exemple suivant :\n"
)

cot_prompt = (
    "Tu es un expert en logique. "
    + "Classifie la relation entre la prémisse et l'hypothèse en utilisant uniquement : "
    + "Conséquence, Contradiction, ou Neutre.\n\n"
    + "Voici des exemples de classification corrects :\n"
    + examples_str
    + "\n Instructions pour la nouvelle tâche :\n"
    + "1. Analysez d'abord la logique étape par étape.\n"
    + "2. Conclus ensuite uniquement par le nom de la relation: Conséquence, Contradiction, ou Neutre.\n"
    + "Réponds maintenant pour l'exemple suivant :\n"
)
print(cot_prompt)

from huggingface_hub import login
from kaggle_secrets import UserSecretsClient

token = UserSecretsClient().get_secret("HUGGING_FACE_HUB_TOKEN")
login(token=token)

import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import bitsandbytes as bnb

print(bnb.__version__)

MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device: ", DEVICE)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"

bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, quantization_config=bnb_config, dtype=DTYPE, device_map="auto"
)
model = torch.compile(model)

def predict_batch(
    premise: list[str], hypotheses: list[str], system_prompt: str, batch_size=32
) -> list[str]:
    """
    Génère les prédictions à partir des prémises et hypothèses.
    """
    # Formatage des prompts
    prompts = [
        tokenizer.apply_chat_template(
            [
                {
                    "role": "user",
                    "content": f"{system_prompt}\nPrémisse : {p}\nHypothèse : {h}\nClassification :",
                }
            ],
            tokenize=False,
            add_generation_prompt=True,
        )
        for p, h in zip(premise, hypotheses)
    ]

    results = []
    for i in range(0, len(prompts), batch_size):
        batch = prompts[i : i + batch_size]

        inputs = tokenizer(
            batch, return_tensors="pt", padding=True, truncation=True, max_length=1024
        ).to(DEVICE)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                do_sample=False,
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id,
                use_cache=True,
            )

        # On garde uniquement les derniers tokens pour éviter de décoder le prompt à nouveau
        generated_ids = outputs[:, inputs.input_ids.shape[1] :]
        results.extend(tokenizer.batch_decode(generated_ids, skip_special_tokens=True))

    return results


def parse_response(response_col: str) -> pl.Expr:
    # Expression Regex pour trouver le label (insensible à la casse)
    label_pattern = r"(?i)\b(Conséquence|Contradiction|Neutre)\b"

    return (
        pl.col(response_col)
        # 1. Trouver toutes les labels dans une listes
        .str.extract_all(label_pattern)
        # 2. Garder le dernier labels
        .list.last()
        .str.to_titlecase()
        .fill_null("Prédiction incorrecte")
    )

sample_size = df_val.height
df_val = df_val.head(sample_size)

print(f"Taille de l'ensemble d'évaluation: {df_val.height} éléments")
print("Répartition des classes dans l'ensemble d'évaluation:")
print(df_val.get_column("label").value_counts(sort=True))

premises = df_val.get_column("premise").to_list()
hypotheses = df_val.get_column("hypothesis").to_list()
valid_labels = ["Neutre", "Conséquence", "Contradiction"]

# Générer les prédictions
predictions = predict_batch(premises, hypotheses, zero_shot_prompt)

# Ajout et parsing des prédictions
df_val = df_val.with_columns(predictions_zero_shot=pl.Series(predictions)).with_columns(
    parse_response("predictions_zero_shot")
)

# Prédictions non valides
valid = pl.col("predictions_zero_shot").is_in(valid_labels)

print(
    f"Nombre de prédictions invalides: ",
    df_val.filter(~valid).height,
    "/",
    df_val.height,
)
df_val.filter(~valid)

from sklearn.metrics import classification_report
import numpy as np

print(
    classification_report(
        df_val.get_column("label"),
        df_val.get_column("predictions_zero_shot"),
        zero_division=np.nan,
    )
)

# Générer les prédictions
predictions = predict_batch(premises, hypotheses, few_shot_prompt)

# Ajout et parsing des prédictions
df_val = df_val.with_columns(predictions_few_shot=pl.Series(predictions)).with_columns(
    parse_response("predictions_few_shot")
)

# Prédictions non valides
valid = pl.col("predictions_few_shot").is_in(valid_labels)

print(
    f"Nombre de prédictions invalides: ",
    df_val.filter(~valid).height,
    "/",
    df_val.height,
)
df_val.filter(~valid)

print(
    classification_report(
        df_val.get_column("label"),
        df_val.get_column("predictions_few_shot"),
        zero_division=np.nan,
    )
)

# Générer les prédictions
predictions = predict_batch(premises, hypotheses, cot_prompt)

# Ajout et parsing des prédictions
df_val = df_val.with_columns(
    predictions_cot_raisonnement=pl.Series(predictions)
).with_columns(predictions_cot=parse_response("predictions_cot_raisonnement"))

# Prédictions non valides
valid = pl.col("predictions_cot").is_in(valid_labels)
invalid_examples = df_val.filter(~valid)

print(f"Nombre de prédictions invalides: ", invalid_examples.height, "/", df_val.height)
print("Exemple de prédictions invalides:\n")
for row in invalid_examples.sample(2).to_dicts():
    print("-" * 50)
    print(f"**Premise**: {row['premise']}")
    print(f"**Hypothesis**: {row['hypothesis']}")
    print(f"**Label**: {row['label']}")
    print(f"\n**Reasoning (CoT)**:\n{row['predictions_cot_raisonnement']}")
    print(f"\n**Final Prediction**: {row['predictions_cot']}")
    print("-" * 50 + "\n")

# Exemple de raisonnement sans parsing possible
df_val.filter(~valid).get_column("predictions_cot_raisonnement").item(0)

print(
    classification_report(
        df_val.get_column("label"),
        df_val.get_column("predictions_cot"),
        zero_division=np.nan,
    )
)

df_val.sample(10)