Bug fixes

src/data/make_dataset.py

@@ -10,5 +10,6 @@ def make_dataset(dataset='cnn_dailymail', split='train', version="3.0.0"):
     df['output_text'] = dataset['target']
     return df
 
+
 if __name__ == '__main__':
-    make_dataset(dataset='cnn_dailymail', split='train', version="3.0.0")
+    make_dataset(dataset='cnn_dailymail', split='train', version="3.0.0")

src/models/model.py

@@ -302,16 +302,7 @@ class Summarization:
             tokenizer=self.tokenizer, model=self.model, output=outputdir
         )
 
-        # checkpoint_callback = ModelCheckpoint(
-        #     dirpath="checkpoints",
-        #     filename="best-checkpoint-{epoch}-{train_loss:.2f}",
-        #     save_top_k=-1,
-        #     verbose=True,
-        #     monitor="train_loss",
-        #     mode="min",
-        # )
-
-        logger = MLFlowLogger(experiment_name="Summarization")
+        logger = MLFlowLogger(experiment_name="Summarization",tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow")
 
         early_stop_callback = (
             [

src/models/predict_model.py

@@ -1,4 +1,5 @@
 from .model import Summarization
+from .make_dataset import make_dataset
 
 def predict_model(text):
     """
@@ -8,4 +9,9 @@ def predict_model(text):
     model.load_model()
     pre_summary = model.predict(text)
     return pre_summary
-    
+
+    
+if __name__ == '__main__':
+    text = make_dataset(split="test")['input_text']
+    pre_summary = predict_model(text)
+    print(pre_summary)  

src/models/train_model.py

@@ -7,9 +7,12 @@ def train_model():
     """
     # Load the data
     train_df = make_dataset(split = 'train')
-    eval_df = make_dataset(split = 'test')
+    eval_df = make_dataset(split = 'val')
 
     model = Summarization()
     model.from_pretrained('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.save_model()
+
+if __name__ == '__main__':
+    train_model()