Spaces:
Sleeping
Sleeping
File size: 4,233 Bytes
0cebe35 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
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()
|