black style added

Makefile

@@ -35,6 +35,7 @@ clean:
 ## Lint using flake8
 lint:
 	flake8 src
+	black src
 
 ## Upload Data to default DVC remote
 push:

src/data/make_dataset.py

@@ -5,22 +5,21 @@ import os
 import pprint
 
 
-
-def make_dataset(dataset='cnn_dailymail', split='train'):
+def make_dataset(dataset="cnn_dailymail", split="train"):
     """make dataset for summarisation"""
-    if not os.path.exists('data/raw'):
-        os.makedirs('data/raw')
-    dataset = load_dataset(dataset, '3.0.0', split=split)
+    if not os.path.exists("data/raw"):
+        os.makedirs("data/raw")
+    dataset = load_dataset(dataset, "3.0.0", split=split)
     df = pd.DataFrame()
-    df['article'] = dataset['article']
-    df['highlights'] = dataset['highlights']
-    df.to_csv('data/raw/{}.csv'.format(split))
+    df["article"] = dataset["article"]
+    df["highlights"] = dataset["highlights"]
+    df.to_csv("data/raw/{}.csv".format(split))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     with open("params.yml") as f:
         params = yaml.safe_load(f)
     pprint.pprint(params)
-    make_dataset(dataset=params['data'], split='train')
-    make_dataset(dataset=params['data'], split='test')
-    make_dataset(dataset=params['data'], split='validation')
+    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

@@ -3,20 +3,20 @@ import yaml
 import os
 
 
-def process_data(split='train'):
+def process_data(split="train"):
 
     with open("params.yml") as f:
         params = yaml.safe_load(f)
 
-    df = pd.read_csv('data/raw/{}.csv'.format(split))
-    df.columns = ['Unnamed: 0', 'input_text', 'output_text']
-    df = df.sample(frac=params['split'], replace=True, random_state=1)
+    df = pd.read_csv("data/raw/{}.csv".format(split))
+    df.columns = ["Unnamed: 0", "input_text", "output_text"]
+    df = df.sample(frac=params["split"], replace=True, random_state=1)
     if os.path.exists("data/raw/{}.csv".format(split)):
         os.remove("data/raw/{}.csv".format(split))
-    df.to_csv('data/processed/{}.csv'.format(split))
+    df.to_csv("data/processed/{}.csv".format(split))
 
 
-if __name__ == '__main__':
-    process_data(split='train')
-    process_data(split='test')
-    process_data(split='validation')
+if __name__ == "__main__":
+    process_data(split="train")
+    process_data(split="test")
+    process_data(split="validation")

src/models/__init__.py

@@ -1 +1 @@
-from .model import Summarization
+from .model import Summarization

src/models/evaluate_model.py

@@ -13,14 +13,14 @@ def evaluate_model():
     with open("params.yml") as f:
         params = yaml.safe_load(f)
 
-    test_df = pd.read_csv('data/processed/test.csv')[:25]
+    test_df = pd.read_csv("data/processed/test.csv")[:25]
     model = Summarization()
-    model.load_model(model_type=params['model_type'], model_dir=params['model_dir'])
-    results = model.evaluate(test_df=test_df, metrics=params['metric'])
+    model.load_model(model_type=params["model_type"], model_dir=params["model_dir"])
+    results = model.evaluate(test_df=test_df, metrics=params["metric"])
 
-    with open('reports/metrics.txt', 'w') as fp:
+    with open("reports/metrics.txt", "w") as fp:
         json.dump(results, fp)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     evaluate_model()

src/models/model.py

@@ -3,7 +3,10 @@ import pandas as pd
 from transformers import (
     AdamW,
     T5ForConditionalGeneration,
-    T5TokenizerFast as T5Tokenizer, MT5Tokenizer, MT5ForConditionalGeneration, ByT5Tokenizer,
+    T5TokenizerFast as T5Tokenizer,
+    MT5Tokenizer,
+    MT5ForConditionalGeneration,
+    ByT5Tokenizer,
 )
 from torch.utils.data import Dataset, DataLoader
 import pytorch_lightning as pl
@@ -28,11 +31,11 @@ class DataModule(Dataset):
     """
 
     def __init__(
-            self,
-            data: pd.DataFrame,
-            tokenizer: T5Tokenizer,
-            source_max_token_len: int = 512,
-            target_max_token_len: int = 512,
+        self,
+        data: pd.DataFrame,
+        tokenizer: T5Tokenizer,
+        source_max_token_len: int = 512,
+        target_max_token_len: int = 512,
     ):
         """
         :param data:
@@ -72,9 +75,7 @@ class DataModule(Dataset):
         )
 
         labels = output_encoding["input_ids"]
-        labels[
-            labels == 0
-            ] = -100
+        labels[labels == 0] = -100
 
         return dict(
             keywords=data_row["input_text"],
@@ -88,15 +89,15 @@ class DataModule(Dataset):
 
 class PLDataModule(LightningDataModule):
     def __init__(
-            self,
-            train_df: pd.DataFrame,
-            test_df: pd.DataFrame,
-            tokenizer: T5Tokenizer,
-            source_max_token_len: int = 512,
-            target_max_token_len: int = 512,
-            batch_size: int = 4,
-            split: float = 0.1,
-            num_workers: int = 2
+        self,
+        train_df: pd.DataFrame,
+        test_df: pd.DataFrame,
+        tokenizer: T5Tokenizer,
+        source_max_token_len: int = 512,
+        target_max_token_len: int = 512,
+        batch_size: int = 4,
+        split: float = 0.1,
+        num_workers: int = 2,
     ):
         """
         :param data_df:
@@ -131,28 +132,45 @@ class PLDataModule(LightningDataModule):
         )
 
     def train_dataloader(self):
-        """ training dataloader """
+        """training dataloader"""
         return DataLoader(
-            self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
         )
 
     def test_dataloader(self):
-        """ test dataloader """
+        """test dataloader"""
         return DataLoader(
-            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
         )
 
     def val_dataloader(self):
-        """ validation dataloader """
+        """validation dataloader"""
         return DataLoader(
-            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
         )
 
 
 class LightningModel(LightningModule):
-    """ PyTorch Lightning Model class"""
+    """PyTorch Lightning Model class"""
 
-    def __init__(self, tokenizer, model, learning_rate, adam_epsilon, weight_decay, output: str = "outputs"):
+    def __init__(
+        self,
+        tokenizer,
+        model,
+        learning_rate,
+        adam_epsilon,
+        weight_decay,
+        output: str = "outputs",
+    ):
         """
         initiates a PyTorch Lightning Model
         Args:
@@ -169,7 +187,7 @@ class LightningModel(LightningModule):
         self.weight_decay = weight_decay
 
     def forward(self, input_ids, attention_mask, decoder_attention_mask, labels=None):
-        """ forward step """
+        """forward step"""
         output = self.model(
             input_ids,
             attention_mask=attention_mask,
@@ -180,7 +198,7 @@ class LightningModel(LightningModule):
         return output.loss, output.logits
 
     def training_step(self, batch, batch_size):
-        """ training step """
+        """training step"""
         input_ids = batch["keywords_input_ids"]
         attention_mask = batch["keywords_attention_mask"]
         labels = batch["labels"]
@@ -196,7 +214,7 @@ class LightningModel(LightningModule):
         return loss
 
     def validation_step(self, batch, batch_size):
-        """ validation step """
+        """validation step"""
         input_ids = batch["keywords_input_ids"]
         attention_mask = batch["keywords_attention_mask"]
         labels = batch["labels"]
@@ -212,7 +230,7 @@ class LightningModel(LightningModule):
         return loss
 
     def test_step(self, batch, batch_size):
-        """ test step """
+        """test step"""
         input_ids = batch["keywords_input_ids"]
         attention_mask = batch["keywords_attention_mask"]
         labels = batch["labels"]
@@ -229,29 +247,39 @@ class LightningModel(LightningModule):
         return loss
 
     def configure_optimizers(self):
-        """ configure optimizers """
+        """configure optimizers"""
         model = self.model
         no_decay = ["bias", "LayerNorm.weight"]
         optimizer_grouped_parameters = [
             {
-                "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
+                "params": [
+                    p
+                    for n, p in model.named_parameters()
+                    if not any(nd in n for nd in no_decay)
+                ],
                 "weight_decay": self.weight_decay,
             },
             {
-                "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
+                "params": [
+                    p
+                    for n, p in model.named_parameters()
+                    if any(nd in n for nd in no_decay)
+                ],
                 "weight_decay": 0.0,
             },
         ]
-        optimizer = AdamW(optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon)
+        optimizer = AdamW(
+            optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon
+        )
         self.opt = optimizer
         return [optimizer]
 
 
 class Summarization:
-    """ Custom Summarization class """
+    """Custom Summarization class"""
 
     def __init__(self) -> None:
-        """ initiates Summarization class """
+        """initiates Summarization class"""
         pass
 
     def from_pretrained(self, model_type="t5", model_name="t5-base") -> None:
@@ -278,20 +306,20 @@ class Summarization:
             )
 
     def train(
-            self,
-            train_df: pd.DataFrame,
-            eval_df: pd.DataFrame,
-            source_max_token_len: int = 512,
-            target_max_token_len: int = 512,
-            batch_size: int = 8,
-            max_epochs: int = 5,
-            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,
-            weight_decay: float = 0.0001
+        self,
+        train_df: pd.DataFrame,
+        eval_df: pd.DataFrame,
+        source_max_token_len: int = 512,
+        target_max_token_len: int = 512,
+        batch_size: int = 8,
+        max_epochs: int = 5,
+        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,
+        weight_decay: float = 0.0001,
     ):
         """
         trains T5/MT5 model on custom dataset
@@ -323,12 +351,18 @@ class Summarization:
         )
 
         self.T5Model = LightningModel(
-            tokenizer=self.tokenizer, model=self.model, output=outputdir,
-            learning_rate=learning_rate, adam_epsilon=adam_epsilon, weight_decay=weight_decay
+            tokenizer=self.tokenizer,
+            model=self.model,
+            output=outputdir,
+            learning_rate=learning_rate,
+            adam_epsilon=adam_epsilon,
+            weight_decay=weight_decay,
         )
 
-        MLlogger = MLFlowLogger(experiment_name="Summarization",
-                                tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow")
+        MLlogger = MLFlowLogger(
+            experiment_name="Summarization",
+            tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow",
+        )
 
         WandLogger = WandbLogger(project="summarization-dagshub")
 
@@ -361,7 +395,7 @@ class Summarization:
         trainer.fit(self.T5Model, self.data_module)
 
     def load_model(
-            self, model_type: str = 't5', model_dir: str = "models", use_gpu: bool = False
+        self, model_type: str = "t5", model_dir: str = "models", use_gpu: bool = False
     ):
         """
         loads a checkpoint for inferencing/prediction
@@ -390,16 +424,15 @@ class Summarization:
             if torch.cuda.is_available():
                 self.device = torch.device("cuda")
             else:
-                raise Exception("exception ---> no gpu found. set use_gpu=False, to use CPU")
+                raise Exception(
+                    "exception ---> no gpu found. set use_gpu=False, to use CPU"
+                )
         else:
             self.device = torch.device("cpu")
 
         self.model = self.model.to(self.device)
 
-    def save_model(
-            self,
-            model_dir="models"
-    ):
+    def save_model(self, model_dir="models"):
         """
         Save model to dir
         :param model_dir:
@@ -410,19 +443,19 @@ class Summarization:
         self.model.save_pretrained(path)
 
     def predict(
-            self,
-            source_text: str,
-            max_length: int = 512,
-            num_return_sequences: int = 1,
-            num_beams: int = 2,
-            top_k: int = 50,
-            top_p: float = 0.95,
-            do_sample: bool = True,
-            repetition_penalty: float = 2.5,
-            length_penalty: float = 1.0,
-            early_stopping: bool = True,
-            skip_special_tokens: bool = True,
-            clean_up_tokenization_spaces: bool = True,
+        self,
+        source_text: str,
+        max_length: int = 512,
+        num_return_sequences: int = 1,
+        num_beams: int = 2,
+        top_k: int = 50,
+        top_p: float = 0.95,
+        do_sample: bool = True,
+        repetition_penalty: float = 2.5,
+        length_penalty: float = 1.0,
+        early_stopping: bool = True,
+        skip_special_tokens: bool = True,
+        clean_up_tokenization_spaces: bool = True,
     ):
         """
         generates prediction for T5/MT5 model
@@ -465,14 +498,10 @@ class Summarization:
         )
         return preds
 
-    def evaluate(
-            self,
-            test_df: pd.DataFrame,
-            metrics: str = "rouge"
-    ):
+    def evaluate(self, test_df: pd.DataFrame, metrics: str = "rouge"):
         metric = load_metric(metrics)
-        input_text = test_df['input_text']
-        references = test_df['output_text']
+        input_text = test_df["input_text"]
+        references = test_df["output_text"]
         references = references.to_list()
 
         predictions = [self.predict(x) for x in tqdm(input_text)]
@@ -480,49 +509,49 @@ class Summarization:
         results = metric.compute(predictions=predictions, references=references)
 
         output = {
-            'Rouge 1': {
-                'Rouge_1 Low Precision': results["rouge1"].low.precision,
-                'Rouge_1 Low recall': results["rouge1"].low.recall,
-                'Rouge_1 Low F1': results["rouge1"].low.fmeasure,
-                'Rouge_1 Mid Precision': results["rouge1"].mid.precision,
-                'Rouge_1 Mid recall': results["rouge1"].mid.recall,
-                'Rouge_1 Mid F1': results["rouge1"].mid.fmeasure,
-                'Rouge_1 High Precision': results["rouge1"].high.precision,
-                'Rouge_1 High recall': results["rouge1"].high.recall,
-                'Rouge_1 High F1': results["rouge1"].high.fmeasure,
+            "Rouge 1": {
+                "Rouge_1 Low Precision": results["rouge1"].low.precision,
+                "Rouge_1 Low recall": results["rouge1"].low.recall,
+                "Rouge_1 Low F1": results["rouge1"].low.fmeasure,
+                "Rouge_1 Mid Precision": results["rouge1"].mid.precision,
+                "Rouge_1 Mid recall": results["rouge1"].mid.recall,
+                "Rouge_1 Mid F1": results["rouge1"].mid.fmeasure,
+                "Rouge_1 High Precision": results["rouge1"].high.precision,
+                "Rouge_1 High recall": results["rouge1"].high.recall,
+                "Rouge_1 High F1": results["rouge1"].high.fmeasure,
+            },
+            "Rouge 2": {
+                "Rouge_2 Low Precision": results["rouge2"].low.precision,
+                "Rouge_2 Low recall": results["rouge2"].low.recall,
+                "Rouge_2 Low F1": results["rouge2"].low.fmeasure,
+                "Rouge_2 Mid Precision": results["rouge2"].mid.precision,
+                "Rouge_2 Mid recall": results["rouge2"].mid.recall,
+                "Rouge_2 Mid F1": results["rouge2"].mid.fmeasure,
+                "Rouge_2 High Precision": results["rouge2"].high.precision,
+                "Rouge_2 High recall": results["rouge2"].high.recall,
+                "Rouge_2 High F1": results["rouge2"].high.fmeasure,
             },
-            'Rouge 2': {
-                'Rouge_2 Low Precision': results["rouge2"].low.precision,
-                'Rouge_2 Low recall': results["rouge2"].low.recall,
-                'Rouge_2 Low F1': results["rouge2"].low.fmeasure,
-                'Rouge_2 Mid Precision': results["rouge2"].mid.precision,
-                'Rouge_2 Mid recall': results["rouge2"].mid.recall,
-                'Rouge_2 Mid F1': results["rouge2"].mid.fmeasure,
-                'Rouge_2 High Precision': results["rouge2"].high.precision,
-                'Rouge_2 High recall': results["rouge2"].high.recall,
-                'Rouge_2 High F1': results["rouge2"].high.fmeasure,
+            "Rouge L": {
+                "Rouge_L Low Precision": results["rougeL"].low.precision,
+                "Rouge_L Low recall": results["rougeL"].low.recall,
+                "Rouge_L Low F1": results["rougeL"].low.fmeasure,
+                "Rouge_L Mid Precision": results["rougeL"].mid.precision,
+                "Rouge_L Mid recall": results["rougeL"].mid.recall,
+                "Rouge_L Mid F1": results["rougeL"].mid.fmeasure,
+                "Rouge_L High Precision": results["rougeL"].high.precision,
+                "Rouge_L High recall": results["rougeL"].high.recall,
+                "Rouge_L High F1": results["rougeL"].high.fmeasure,
             },
-            'Rouge L': {
-                'Rouge_L Low Precision': results["rougeL"].low.precision,
-                'Rouge_L Low recall': results["rougeL"].low.recall,
-                'Rouge_L Low F1': results["rougeL"].low.fmeasure,
-                'Rouge_L Mid Precision': results["rougeL"].mid.precision,
-                'Rouge_L Mid recall': results["rougeL"].mid.recall,
-                'Rouge_L Mid F1': results["rougeL"].mid.fmeasure,
-                'Rouge_L High Precision': results["rougeL"].high.precision,
-                'Rouge_L High recall': results["rougeL"].high.recall,
-                'Rouge_L High F1': results["rougeL"].high.fmeasure,
+            "rougeLsum": {
+                "rougeLsum Low Precision": results["rougeLsum"].low.precision,
+                "rougeLsum Low recall": results["rougeLsum"].low.recall,
+                "rougeLsum Low F1": results["rougeLsum"].low.fmeasure,
+                "rougeLsum Mid Precision": results["rougeLsum"].mid.precision,
+                "rougeLsum Mid recall": results["rougeLsum"].mid.recall,
+                "rougeLsum Mid F1": results["rougeLsum"].mid.fmeasure,
+                "rougeLsum High Precision": results["rougeLsum"].high.precision,
+                "rougeLsum High recall": results["rougeLsum"].high.recall,
+                "rougeLsum High F1": results["rougeLsum"].high.fmeasure,
             },
-            'rougeLsum': {
-                'rougeLsum Low Precision': results["rougeLsum"].low.precision,
-                'rougeLsum Low recall': results["rougeLsum"].low.recall,
-                'rougeLsum Low F1': results["rougeLsum"].low.fmeasure,
-                'rougeLsum Mid Precision': results["rougeLsum"].mid.precision,
-                'rougeLsum Mid recall': results["rougeLsum"].mid.recall,
-                'rougeLsum Mid F1': results["rougeLsum"].mid.fmeasure,
-                'rougeLsum High Precision': results["rougeLsum"].high.precision,
-                'rougeLsum High recall': results["rougeLsum"].high.recall,
-                'rougeLsum High F1': results["rougeLsum"].high.fmeasure,
-            }
         }
         return output

src/models/predict_model.py

@@ -11,14 +11,13 @@ def predict_model(text):
     with open("params.yml") as f:
         params = yaml.safe_load(f)
 
-
     model = Summarization()
-    model.load_model(model_type=params['model_type'], model_dir=params['model_dir'])
+    model.load_model(model_type=params["model_type"], model_dir=params["model_dir"])
     pre_summary = model.predict(text)
     return pre_summary
 
 
-if __name__ == '__main__':
-    text = pd.load_csv('data/processed/test.csv')['input_text'][0]
+if __name__ == "__main__":
+    text = pd.load_csv("data/processed/test.csv")["input_text"][0]
     pre_summary = predict_model(text)
     print(pre_summary)

src/models/train_model.py

@@ -14,28 +14,35 @@ def train_model():
         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")
 
-    train_df = train_df.sample(frac=params['split'], replace=True, random_state=1)
-    eval_df = eval_df.sample(frac=params['split'], replace=True, random_state=1)
+    train_df = train_df.sample(frac=params["split"], replace=True, random_state=1)
+    eval_df = eval_df.sample(frac=params["split"], replace=True, random_state=1)
 
     model = Summarization()
-    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['epochs'],
-                use_gpu=params['use_gpu'], learning_rate=float(params['learning_rate']),
-                num_workers=int(params['num_workers']))
-
-    model.save_model(model_dir=params['model_dir'])
-
-    with open('wandb/latest-run/files/wandb-summary.json') as json_file:
+    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["epochs"],
+        use_gpu=params["use_gpu"],
+        learning_rate=float(params["learning_rate"]),
+        num_workers=int(params["num_workers"]),
+    )
+
+    model.save_model(model_dir=params["model_dir"])
+
+    with open("wandb/latest-run/files/wandb-summary.json") as json_file:
         data = json.load(json_file)
 
-    with open('reports/training_metrics.txt', 'w') as fp:
+    with open("reports/training_metrics.txt", "w") as fp:
         json.dump(data, fp)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     train_model()