Fixes

src/models/evaluate_model.py

@@ -10,7 +10,7 @@ def evaluate_model():
     test_df =  pd.load_csv('../../data/processed/test.csv')
     model = Summarization()
     model.load_model()
-    results = model.evaluate(test_df=test_df)
+    results = model.evaluate(test_df=test_df,metrics="rouge")
     with dagshub.dagshub_logger() as logger:
         logger.log_metrics(results)
     return results

src/models/model.py

@@ -6,7 +6,7 @@ from dagshub.pytorch_lightning import DAGsHubLogger
 from transformers import (
     AdamW,
     T5ForConditionalGeneration,
-    T5TokenizerFast as T5Tokenizer,
+    T5TokenizerFast as T5Tokenizer, MT5Tokenizer, MT5ForConditionalGeneration,ByT5Tokenizer,
 )
 from torch.utils.data import Dataset, DataLoader
 import pytorch_lightning as pl
@@ -250,16 +250,28 @@ class Summarization:
         """ initiates Summarization class """
         pass
 
-    def from_pretrained(self, model_name="t5-base") -> None:
+    def from_pretrained(self,model_type = "t5", model_name="t5-base") -> None:
         """
         loads T5/MT5 Model model for training/finetuning
         Args:
             model_name (str, optional): exact model architecture name, "t5-base" or "t5-large". Defaults to "t5-base".
+            :param model_type:
         """
-        self.tokenizer = T5Tokenizer.from_pretrained(f"{model_name}")
-        self.model = T5ForConditionalGeneration.from_pretrained(
-            f"{model_name}", return_dict=True
-        )
+        if model_type == "t5":
+            self.tokenizer = T5Tokenizer.from_pretrained(f"{model_name}")
+            self.model = T5ForConditionalGeneration.from_pretrained(
+                f"{model_name}", return_dict=True
+            )
+        elif model_type == "mt5":
+            self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_name}")
+            self.model = MT5ForConditionalGeneration.from_pretrained(
+                f"{model_name}", return_dict=True
+            )
+        elif model_type == "byt5":
+            self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
+            self.model = T5ForConditionalGeneration.from_pretrained(
+                f"{model_name}", return_dict=True
+            )
 
     def train(
             self,

src/models/train_model.py

@@ -1,18 +1,20 @@
 from src.models.model import Summarization
 import pandas as pd
 
+
 def train_model():
     """
     Train the model
     """
     # Load the data
-    train_df = pd.load_csv('../../data/processed/train.csv')
-    eval_df = pd.load_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-base')
+    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()
 
+
 if __name__ == '__main__':
-    train_model()
+    train_model()