| from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling | |
| from datasets import load_dataset | |
| from model import SASOKModel, SASOKConfig | |
| from tokenizer import tokenizer | |
| dataset = load_dataset("wikitext", "wikitext-2-raw-v1") | |
| tokenized = dataset.map(lambda e: tokenizer(e["text"], truncation=True, padding="max_length"), batched=True) | |
| config = SASOKConfig() | |
| model = SASOKModel(config) | |
| training_args = TrainingArguments( | |
| output_dir="./sasok_output", | |
| evaluation_strategy="steps", | |
| eval_steps=500, | |
| per_device_train_batch_size=4, | |
| num_train_epochs=3, | |
| save_steps=1000, | |
| logging_dir="./logs" | |
| ) | |
| data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=tokenized["train"], | |
| eval_dataset=tokenized["validation"], | |
| tokenizer=tokenizer, | |
| data_collator=data_collator | |
| ) | |
| trainer.train() |