Spaces:
Sleeping
Sleeping
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() | |