Updates

.gitignore

@@ -88,3 +88,6 @@ coverage.xml
 
 # Mypy cache
 .mypy_cache/
+
+.idea 
+.vscode

src/data/make_dataset.py

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

src/models/model.py

@@ -340,7 +340,7 @@ class Summarization:
         trainer.fit(self.T5Model, self.data_module)
 
     def load_model(
-            self, model_dir: str = "models", use_gpu: bool = False
+            self, model_dir: str = "../../models", use_gpu: bool = False
     ):
         """
         loads a checkpoint for inferencing/prediction

src/models/predict_model.py

@@ -1,2 +1,11 @@
 from .model import Summarization
 
+def predict_model(text):
+    """
+    Predict the summary of the given text.
+    """
+    model = Summarization()
+    model.load_model()
+    pre_summary = model.predict(text)
+    return pre_summary
+    

src/models/train_model.py

@@ -0,0 +1,15 @@
+from .model import Summarization
+from data.make_dataset import make_dataset
+
+def train_model():
+    """
+    Train the model
+    """
+    # Load the data
+    train_df = make_dataset(split = 'train')
+    eval_df = make_dataset(split = 'test')
+
+    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()