added params

params.yml

@@ -0,0 +1,10 @@
+data: cnn_dailymail
+batch_size: 4
+num_workers: 2
+model_type: t5
+model_name: t5-base
+learning_rate: 1e-4
+epochs: 5
+source_dir: src
+model_dir: models
+metric: rouge

src/data/make_dataset.py

@@ -1,3 +1,4 @@
+import yaml
 from datasets import load_dataset
 import pandas as pd
 
@@ -8,10 +9,13 @@ def make_dataset(dataset='cnn_dailymail', split='train'):
     df = pd.DataFrame()
     df['article'] = dataset['article']
     df['highlights'] = dataset['highlights']
-    df.to_csv('C:/Users/gbhat/Documents/GitHub/summarization/data/raw/{}.csv'.format(split))
+    df.to_csv('data/raw/{}.csv'.format(split))
 
 
 if __name__ == '__main__':
-    make_dataset(dataset='cnn_dailymail', split='train')
-    make_dataset(dataset='cnn_dailymail', split='test')
-    make_dataset(dataset='cnn_dailymail', split='validation')
+    with open("params.yml") as f:
+        params = yaml.safe_load(f)
+
+    make_dataset(dataset=params['data'], split='train')
+    make_dataset(dataset=params['data'], split='test')
+    make_dataset(dataset=params['data'], split='validation')

src/data/process_data.py

@@ -2,10 +2,10 @@ import pandas as pd
 
 
 def process_data(split='train'):
-    df = pd.read_csv('C:/Users/gbhat/Documents/GitHub/summarization/data/raw/{}.csv'.format(split))
+
+    df = pd.read_csv('data/raw/{}.csv'.format(split))
     df.columns = ['Unnamed: 0', 'input_text', 'output_text']
-    print(df.columns)
-    df.to_csv('C:/Users/gbhat/Documents/GitHub/summarization/data/processed/{}.csv'.format(split))
+    df.to_csv('data/processed/{}.csv'.format(split))
 
 
 if __name__ == '__main__':

src/models/evaluate_model.py

@@ -1,16 +1,22 @@
 import dagshub
+import yaml
 
 from src.models.model import Summarization
 import pandas as pd
 
+
 def evaluate_model():
     """
     Evaluate model using rouge measure
     """
-    test_df =  pd.load_csv('../../data/processed/test.csv')
+    with open("params.yml") as f:
+        params = yaml.safe_load(f)
+
+    test_df = pd.load_csv('data/processed/test.csv')
     model = Summarization()
-    model.load_model()
-    results = model.evaluate(test_df=test_df,metrics="rouge")
+    model.load_model(model_dir=params['model_dir'])
+    results = model.evaluate(test_df=test_df, metrics=params['metric'])
+
     with dagshub.dagshub_logger() as logger:
         logger.log_metrics(results)
     return results

src/models/model.py

@@ -94,7 +94,8 @@ class PLDataModule(LightningDataModule):
             source_max_token_len: int = 512,
             target_max_token_len: int = 512,
             batch_size: int = 4,
-            split: float = 0.1
+            split: float = 0.1,
+            num_workers: int = 2
     ):
         """
         :param data_df:
@@ -112,6 +113,7 @@ class PLDataModule(LightningDataModule):
         self.target_max_token_len = target_max_token_len
         self.source_max_token_len = source_max_token_len
         self.tokenizer = tokenizer
+        self.num_workers = num_workers
 
     def setup(self, stage=None):
         self.train_dataset = DataModule(
@@ -130,26 +132,26 @@ class PLDataModule(LightningDataModule):
     def train_dataloader(self):
         """ training dataloader """
         return DataLoader(
-            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=2
+            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers
         )
 
     def test_dataloader(self):
         """ test dataloader """
         return DataLoader(
-            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=2
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers
         )
 
     def val_dataloader(self):
         """ validation dataloader """
         return DataLoader(
-            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=2
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers
         )
 
 
 class LightningModel(LightningModule):
     """ PyTorch Lightning Model class"""
 
-    def __init__(self, tokenizer, model, output: str = "outputs"):
+    def __init__(self, tokenizer, model, learning_rate, adam_epsilon, output: str = "outputs"):
         """
         initiates a PyTorch Lightning Model
         Args:
@@ -236,7 +238,7 @@ class LightningModel(LightningModule):
                 "weight_decay": 0.0,
             },
         ]
-        optimizer = AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)
+        optimizer = AdamW(optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon)
         self.opt = optimizer
         return [optimizer]
 
@@ -282,6 +284,9 @@ class Summarization:
             use_gpu: bool = True,
             outputdir: str = "models",
             early_stopping_patience_epochs: int = 0,  # 0 to disable early stopping feature
+            learning_rate: float = 0.0001,
+            adam_epsilon: float = 0.01,
+            num_workers: int = 2
     ):
         """
         trains T5/MT5 model on custom dataset
@@ -298,6 +303,8 @@ class Summarization:
             early_stopping_patience_epochs (int, optional): monitors val_loss on epoch end and stops training,
             if val_loss does not improve after the specied number of epochs. set 0 to disable early stopping.
             Defaults to 0 (disabled)
+            :param learning_rate:
+            :param adam_epsilon:
         """
         self.target_max_token_len = target_max_token_len
         self.data_module = PLDataModule(
@@ -307,10 +314,12 @@ class Summarization:
             batch_size=batch_size,
             source_max_token_len=source_max_token_len,
             target_max_token_len=target_max_token_len,
+            num_workers=num_workers,
         )
 
         self.T5Model = LightningModel(
-            tokenizer=self.tokenizer, model=self.model, output=outputdir
+            tokenizer=self.tokenizer, model=self.model, output=outputdir,
+            learning_rate=learning_rate,adam_epsilon=adam_epsilon
         )
 
         MLlogger = MLFlowLogger(experiment_name="Summarization",

src/models/predict_model.py

@@ -1,5 +1,6 @@
 from src.data.make_dataset import make_dataset
 from .model import Summarization
+import pandas as pd
 
 def predict_model(text):
     """
@@ -12,6 +13,6 @@ def predict_model(text):
 
     
 if __name__ == '__main__':
-    text = make_dataset(split="test")['input_text'][0]
+    text = pd.load_csv('data/processed/test.csv')['input_text'][0]
     pre_summary = predict_model(text)
     print(pre_summary)  

src/models/train_model.py

@@ -1,3 +1,5 @@
+import yaml
+
 from src.models.model import Summarization
 import pandas as pd
 
@@ -6,14 +8,22 @@ def train_model():
     """
     Train the model
     """
+    with open("params.yml") as f:
+        params = yaml.safe_load(f)
+
     # Load the data
-    train_df = pd.read_csv('../../data/processed/train.csv')
-    eval_df = pd.read_csv('../../data/processed/validation.csv')
+    train_df = pd.read_csv('data/processed/train.csv')
+    eval_df = pd.read_csv('data/processed/validation.csv')
 
     model = Summarization()
-    model.from_pretrained('t5','t5-base')
-    model.train(train_df=train_df, eval_df=eval_df, batch_size=4, max_epochs=3, use_gpu=True)
-    model.save_model()
+    model.from_pretrained(model_type=params['model_type'], model_name=params['model_name'])
+
+    model.train(train_df=train_df, eval_df=eval_df,
+                batch_size=params['batch_size'], max_epochs=params['max_epoch'],
+                use_gpu=params['use_gpu'], learning_rate=params['learning_rate'],
+                num_workers=params['num_workers'])
+
+    model.save_model(model_dir=params['model_dir'])
 
 
 if __name__ == '__main__':