fixes

dvc.yaml

@@ -8,10 +8,13 @@ stages:
     outs:
       - data/processed/test.csv:
           persist: true
+          cache: true
       - data/processed/train.csv:
           persist: true
+          cache: true
       - data/processed/validation.csv:
           persist: true
+          cache: true
   train:
     cmd: python src/models/train_model.py
     deps:

src/models/model.py

@@ -361,9 +361,6 @@ class Summarization:
 
         trainer.fit(self.T5Model, self.data_module)
 
-        df = pd.read_json(r'wandb/latest-run/files/wandb-summary.json')
-        df.to_csv(r'reports/training_metrics.txt', index=False)
-
     def load_model(
             self, model_type: str = 't5', model_dir: str = "models", use_gpu: bool = False
     ):

src/models/train_model.py

@@ -28,6 +28,9 @@ def train_model():
 
     model.save_model(model_dir=params['model_dir'])
 
+    df = pd.read_json(r'wandb/latest-run/files/wandb-summary.json')
+    df.to_csv(r'reports/training_metrics.txt', index=False)
+
 
 if __name__ == '__main__':
     train_model()