Update train_model.py

src/models/train_model.py

@@ -16,9 +16,9 @@ def train_model():
     eval_df = pd.read_csv('data/processed/validation.csv')
 
     model = Summarization()
-    model.from_pretrained('t5','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.from_pretrained(model_type=params['model_type'], model_name=params['model_name'])
+
+    model.train(train_df=train_df, eval_df=eval_df,
 
 
 if __name__ == '__main__':