Create src/train.py

src/train.py

@@ -0,0 +1,41 @@
+import torch
+from transformers import Trainer, TrainingArguments
+from model import get_model_and_tokenizer
+from data_loader import get_dataloader
+from utils import load_config, set_seed
+
+def main():
+    config = load_config('configs/model_config.yaml')
+    set_seed(config['training']['seed'])
+
+    model, tokenizer = get_model_and_tokenizer(config)
+    train_dataloader = get_dataloader(config, tokenizer, 'train')
+    val_dataloader = get_dataloader(config, tokenizer, 'validation')
+
+    training_args = TrainingArguments(
+        output_dir="./results",
+        num_train_epochs=config['training']['num_epochs'],
+        per_device_train_batch_size=config['training']['batch_size'],
+        per_device_eval_batch_size=config['training']['batch_size'],
+        warmup_steps=500,
+        weight_decay=0.01,
+        logging_dir='./logs',
+        logging_steps=100,
+        evaluation_strategy="steps",
+        eval_steps=1000,
+        save_steps=config['training']['save_every'],
+        load_best_model_at_end=True,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataloader.dataset,
+        eval_dataset=val_dataloader.dataset,
+    )
+
+    trainer.train()
+    trainer.save_model("./final_model")
+
+if __name__ == "__main__":
+    main()