added params

requirements.txt

@@ -6,6 +6,7 @@ torch==1.9.0+cu111
 dagshub==0.1.6
 pandas==1.2.4
 rouge_score
+yaml
 
 # external requirements
 click

src/models/evaluate_model.py

@@ -14,7 +14,7 @@ def evaluate_model():
 
     test_df = pd.load_csv('data/processed/test.csv')
     model = Summarization()
-    model.load_model(model_dir=params['model_dir'])
+    model.load_model(model_type=params['model_type'], model_dir=params['model_dir'])
     results = model.evaluate(test_df=test_df, metrics=params['metric'])
 
     with dagshub.dagshub_logger() as logger:

src/models/predict_model.py

@@ -1,3 +1,5 @@
+import yaml
+
 from .model import Summarization
 import pandas as pd
 
@@ -6,8 +8,12 @@ def predict_model(text):
     """
     Predict the summary of the given text.
     """
+    with open("params.yml") as f:
+        params = yaml.safe_load(f)
+
+
     model = Summarization()
-    model.load_model()
+    model.load_model(model_type=params['model_type'], model_dir=params['model_dir'])
     pre_summary = model.predict(text)
     return pre_summary