training.py

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments
from datasets import Dataset
from datasets import load_dataset
import torch
import numpy as np
from evaluate import load

# Charger le modèle et le tokenizer
model_name = "moussaKam/barthez"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

import pandas as pd
D = pd.read_csv("data/fr_train.csv")
D = D.drop("topic", axis=1)
D = D.drop("url", axis=1)
D = D.drop("title", axis=1)
D = D.drop("date", axis=1)
train_df = D.rename(columns={'text': 'texte', 'summary': 'resume'})
train_df = train_df.head(10)

D = pd.read_csv("data/fr_validation.csv")
D = D.drop("topic", axis=1)
D = D.drop("url", axis=1)
D = D.drop("title", axis=1)
D = D.drop("date", axis=1)
val_df = D.rename(columns={'text': 'texte', 'summary': 'resume'})
val_df = val_df.head(10)

# CORRECTION : Convertir les DataFrames pandas en Dataset HuggingFace
train_dataset = Dataset.from_pandas(train_df)
val_dataset = Dataset.from_pandas(val_df)



# Fonction de prétraitement
def preprocess_function(examples):
    """
    Tokenise les textes et les résumés
    """
    # Tokeniser les textes d'entrée
    inputs = tokenizer(
        examples["texte"],
        max_length=1024,
        truncation=True,
        padding="max_length"
    )
    
    # Tokeniser les résumés cibles
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(
            examples["resume"],
            max_length=150,
            truncation=True,
            padding="max_length"
        )
    
    inputs["labels"] = labels["input_ids"]
    return inputs

# Appliquer le prétraitement
tokenized_train = train_dataset.map(preprocess_function, batched=True)
tokenized_val = val_dataset.map(preprocess_function, batched=True)

# Charger la métrique ROUGE pour l'évaluation (EXPLIC)
rouge = load("rouge")

def compute_metrics(eval_pred):
    """
    Calcule les métriques ROUGE pour évaluer la qualité des résumés
    """
    predictions, labels = eval_pred
    
    # Décoder les prédictions
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    
    # Remplacer -100 dans les labels (padding)
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    # Calculer ROUGE
    result = rouge.compute(
        predictions=decoded_preds,
        references=decoded_labels,
        use_stemmer=True
    )
    
    return {
        "rouge1": result["rouge1"],
        "rouge2": result["rouge2"],
        "rougeL": result["rougeL"]
    }

# Définir les arguments d'entraînement (EXPLIC)
training_args = Seq2SeqTrainingArguments(
    output_dir="./barthez-finetuned-resume",
    eval_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=3,
    weight_decay=0.01,
    save_total_limit=3,
    predict_with_generate=True,
    fp16=torch.cuda.is_available(),  # Utiliser FP16 si GPU disponible
    push_to_hub=False,
    logging_dir="./logs",
    logging_steps=100,
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="rougeL"
)

# Créer le Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

#Lancer l'entraînement
print("Début du fine-tuning...")
trainer.train()

# Sauvegarder le modèle fine-tuné
model.save_pretrained("./barthez-finetuned-resume-final")
tokenizer.save_pretrained("./barthez-finetuned-resume-final")
print("Modèle sauvegardé dans './barthez-finetuned-resume-final'")


# Fonction pour tester le modèle fine-tuné
def tester_modele(texte):
    """
    Teste le modèle fine-tuné sur un nouveau texte
    """
    inputs = tokenizer(
        texte,
        max_length=1024,
        truncation=True,
        return_tensors="pt"
    ).to(model.device)
    
    with torch.no_grad():
        summary_ids = model.generate(
            inputs["input_ids"],
            max_length=150,
            min_length=40,
            num_beams=4,
            length_penalty=2.0,
            early_stopping=True
        )
    
    resume = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return resume

# Exemple de test après l'entraînement
if __name__ == "__main__":
    texte_test = "Votre texte à résumer..."
    print("\nTest du modèle fine-tuné:")
    print(tester_modele(texte_test))