updates

dvc.yaml

@@ -43,6 +43,7 @@ stages:
     deps:
       - models
       - src/visualization/visualize.py
+      - params.yml
     metrics:
       - reports/visualization_metrics.txt:
           cache: false

params.yml

@@ -9,4 +9,5 @@ source_dir: src
 model_dir: models
 metric: rouge
 split: 0.02
-use_gpu: True
+use_gpu: True
+visualise: True

src/visualization/visualize.py

@@ -1,4 +1,6 @@
 import streamlit as st
+import yaml
+
 from models import predict_model
 
 
@@ -22,4 +24,8 @@ def visualize():
 
 
 if __name__ == "__main__":
-    visualize()
+    with open("params.yml") as f:
+        params = yaml.safe_load(f)
+
+    if params['visualise']:
+        visualize()