Merwan6
Commit initial
0cebe35
raw
history blame
4.23 kB
from datasets import load_dataset, DatasetDict, Dataset
from transformers import (
AutoTokenizer, AutoModelForSequenceClassification,
Trainer, TrainingArguments, DataCollatorWithPadding
)
import numpy as np
from utils import compute_metrics
import os
def load_ag_news():
"""
Charge le jeu de données AG News via Hugging Face datasets.
Returns:
DatasetDict: Contenant les splits train/test.
"""
dataset = load_dataset("ag_news")
return dataset
def get_balanced_subset(dataset_split, n_per_class=1000):
"""
Crée un sous-ensemble équilibré contenant `n_per_class` exemples par classe.
Args:
dataset_split (Dataset): Split de type train ou test.
n_per_class (int): Nombre d'exemples à garder par classe.
Returns:
Dataset: Sous-ensemble équilibré.
"""
subsets = []
for label in range(4):
#Filtrage des exemples correspondant à la classe `label`
filtered = dataset_split.filter(lambda example: example['label'] == label)
#Sélection des n premiers exemples (ou tous s’il y en a moins)
subsets.append(filtered.select(range(min(n_per_class, len(filtered)))))
#Fusionner les sous-ensembles
combined_dict = {
key: sum([subset[key] for subset in subsets], []) for key in subsets[0].features.keys()
}
return Dataset.from_dict(combined_dict)
def preprocess_data(dataset, tokenizer):
"""
Tokenise le jeu de données avec troncature et padding.
Args:
dataset (DatasetDict): Données d'entraînement et de test.
tokenizer (AutoTokenizer): Tokenizer à utiliser.
Returns:
DatasetDict: Données tokenisées.
"""
def preprocess(batch):
return tokenizer(batch["text"], truncation=True, padding=True)
return dataset.map(preprocess, batched=True)
def main():
"""
Lance le fine-tuning du modèle BERT sur AG News et sauvegarde le modèle.
"""
#Création des dossiers de sortie
os.makedirs("../models/fine_tuned_model", exist_ok=True)
os.makedirs("../logs", exist_ok=True)
#Chargement du jeu de données
dataset = load_ag_news()
#Création de sous-ensembles équilibrés (entraînement/test)
train_subset = get_balanced_subset(dataset["train"], n_per_class=3000)
test_subset = get_balanced_subset(dataset["test"], n_per_class=1000)
dataset_small = DatasetDict({
"train": train_subset,
"test": test_subset
})
#Chargement du tokenizer
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
#Prétraitement (tokenisation)
encoded = preprocess_data(dataset_small, tokenizer)
#Préparation des entrées avec padding dynamique
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
#Chargement du modèle BERT pour classification
model = AutoModelForSequenceClassification.from_pretrained(
model_name,
num_labels=4 #AG News contient 4 classes
)
#Configuration de l'entraînement
training_args = TrainingArguments(
output_dir="../models/fine_tuned_model",
eval_strategy="epoch",
save_strategy="epoch",
num_train_epochs=3,
per_device_train_batch_size=32,
per_device_eval_batch_size=32,
load_best_model_at_end=True,
metric_for_best_model="accuracy",
logging_dir="../logs",
seed=42
)
#Définition du trainer Hugging Face
trainer = Trainer(
model=model,
args=training_args,
train_dataset=encoded["train"],
eval_dataset=encoded["test"],
tokenizer=tokenizer,
data_collator=data_collator,
compute_metrics=lambda p: compute_metrics(
np.argmax(p.predictions, axis=1), p.label_ids
)
)
#Lancement de l'entraînement
trainer.train()
#Sauvegarde finale du modèle
trainer.save_model("../models/fine_tuned_model")
print("✅ Modèle sauvegardé dans ../models/fine_tuned_model")
if __name__ == "__main__":
main()