model trainer added

config/config.yaml

@@ -17,4 +17,11 @@ data_validation:
 data_transformation:
   root_dir: artifacts/data_transformation
   data_path: artifacts/data_ingestion/samsum_dataset
-  tokenizer_name: google/pegasus-cnn_dailymail
+  tokenizer_name: google/pegasus-cnn_dailymail
+
+
+
+model_trainer:
+  root_dir: artifacts/model_trainer
+  data_path: artifacts/data_transformation/samsum_dataset
+  model_ckpt: google/pegasus-cnn_dailymail

main.py

@@ -1,6 +1,7 @@
 from textsummarizer.pipeline.stage_01_data_ingestion import DataIngestionPipeline
 from textsummarizer.pipeline.stage_02_data_validation import DataValidationPipeline
 from textsummarizer.pipeline.stage_03_data_transformation import DataTransformationPipeline
+from textsummarizer.pipeline.stage_04_model_trainer import ModelTrainerPipeline
 from textsummarizer.logging import logger
 
 STAGE_NAME = "Data Ingestion stage"
@@ -31,6 +32,18 @@ try:
    data_transformaion = DataTransformationPipeline()
    data_transformaion.main()
    logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+except Exception as e:
+        logger.exception(e)
+        raise e
+     
+     
+     
+STAGE_NAME = "Data Traniner stage"
+try:
+   logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<") 
+   model_tranier = ModelTrainerPipeline()
+   model_tranier.main()
+   logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
 except Exception as e:
         logger.exception(e)
         raise e

params.yaml

@@ -1 +1,11 @@
-key : val
+TrainingArguments:
+  num_train_epochs: 1
+  warmup_steps: 500
+  per_device_train_batch_size: 1
+  weight_decay: 0.01
+  logging_steps: 10
+  evaluation_strategy: steps
+  eval_steps: 500
+  save_steps: 1e6
+  gradient_accumulation_steps: 16
+

src/textsummarizer/config/configuration.py

@@ -2,7 +2,8 @@ from textsummarizer.constants import *
 from textsummarizer.utils.common import read_yaml, create_directories
 from textsummarizer.entity.config_entity import (DataIngestionConfig,
                                                  DataValidationConfig,
-                                                 DataTransformationConfig)
+                                                 DataTransformationConfig,
+                                                 ModelTrainerConfig)
 
 class ConfigurationManager:
     def __init__(
@@ -58,5 +59,30 @@ class ConfigurationManager:
         )
 
         return data_transformation_config
+    
+    
+    def get_model_trainer_config(self) -> ModelTrainerConfig:
+        config = self.config.model_trainer
+        params = self.params.TrainingArguments
+
+        create_directories([config.root_dir])
+        
+        
+        model_trainer_config = ModelTrainerConfig(
+            root_dir  = config.root_dir,
+            data_path = config.data_path,
+            model_ckpt = config.model_ckpt,
+            num_train_epochs =params.num_train_epochs,
+            warmup_steps =params.warmup_steps,
+            per_device_train_batch_size = params.per_device_train_batch_size,
+            weight_decay = params.weight_decay,
+            logging_steps = params.logging_steps,
+            evaluation_strategy =params.evaluation_strategy,
+            eval_steps =params.eval_steps,
+            save_steps =  params.save_steps,
+            gradient_accumulation_steps = params.gradient_accumulation_steps
+        )
+        
+        return model_trainer_config
 
         

src/textsummarizer/conponents/model_trainer.py

@@ -0,0 +1,48 @@
+from transformers import TrainingArguments, Trainer
+from transformers import DataCollatorForSeq2Seq
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from datasets import load_dataset, load_from_disk
+from 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"))

src/textsummarizer/entity/config_entity.py

@@ -23,4 +23,18 @@ class DataTransformationConfig:
     data_path : Path
     tokenizer_name : Path
     
-    
+
+@dataclass(frozen=True)
+class ModelTrainerConfig:
+    root_dir : Path
+    data_path : Path
+    model_ckpt  : Path
+    num_train_epochs : int
+    warmup_steps : int
+    per_device_train_batch_size : int
+    weight_decay : float
+    logging_steps : int
+    evaluation_strategy: str
+    eval_steps: int
+    save_steps: float
+    gradient_accumulation_steps: int

src/textsummarizer/pipeline/stage_04_model_trainer.py

@@ -0,0 +1,12 @@
+from textsummarizer.conponents.model_trainer import ModelTrainer
+from textsummarizer.config.configuration import ConfigurationManager
+
+class ModelTrainerPipeline:
+    def __init__(self):
+        pass
+    
+    def main(self):
+        config = ConfigurationManager()
+        model_trainer_config = config.get_model_trainer_config()
+        model_trainer_config = ModelTrainer(config=model_trainer_config)
+        model_trainer_config.train()