File size: 1,290 Bytes
9873cb7 |
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 |
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, TextDataset, DataCollatorForLanguageModeling
# Charger le modèle et le tokenizer GPT-2
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
# Préparer ton dataset (assure-toi que 'train.txt' existe avec tes données)
def load_dataset(file_path, tokenizer, block_size=128):
return TextDataset(
tokenizer=tokenizer,
file_path=file_path,
block_size=block_size
)
# Charger le dataset
train_dataset = load_dataset("train.txt", tokenizer)
# Préparer les arguments pour l'entraînement
training_args = TrainingArguments(
output_dir="./thought_model",
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=2,
save_steps=500,
save_total_limit=2
)
# Préparer le data collator
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False
)
# Lancer l'entraînement
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=train_dataset
)
trainer.train()
# Sauvegarder le modèle fine-tuné
trainer.save_model("./thought_model")
tokenizer.save_pretrained("./thought_model")
|