Spaces:
Sleeping
Sleeping
File size: 2,152 Bytes
f2492e6 a637525 f2492e6 a637525 f2492e6 |
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 |
from transformers import TrainingArguments, Trainer
from transformers import DataCollatorForSeq2Seq
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from datasets import load_dataset, load_from_disk
from src.textsummarizer.entity.config_entity import ModelTrainerConfig
import torch
import os
class ModelTrainer:
def __init__(self, config : ModelTrainerConfig):
self.config = config
os.environ["WANDB_DISABLED"] = "true"
def train(self):
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(self.config.model_ckpt)
model_pegasus = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_ckpt).to(device)
seq2seq_data_collator = DataCollatorForSeq2Seq(tokenizer, model=model_pegasus)
#loading data
dataset_samsum_pt = load_from_disk(self.config.data_path)
trainer_args = TrainingArguments(
output_dir=self.config.root_dir, num_train_epochs=self.config.num_train_epochs, warmup_steps=self.config.warmup_steps,
per_device_train_batch_size=self.config.per_device_train_batch_size, per_device_eval_batch_size=self.config.per_device_train_batch_size,
weight_decay=self.config.weight_decay, logging_steps=self.config.logging_steps,
evaluation_strategy=self.config.evaluation_strategy, eval_steps=self.config.eval_steps, save_steps=1e6,
gradient_accumulation_steps=self.config.gradient_accumulation_steps,
report_to="none"
)
trainer = Trainer(model=model_pegasus, args=trainer_args,
tokenizer=tokenizer, data_collator=seq2seq_data_collator,
train_dataset=dataset_samsum_pt["train"],
eval_dataset=dataset_samsum_pt["validation"])
trainer.train()
## Save model
model_pegasus.save_pretrained(os.path.join(self.config.root_dir,"pegasus-samsum-model"))
## Save tokenizer
tokenizer.save_pretrained(os.path.join(self.config.root_dir,"tokenizer"))
|