Gagan Bhatia commited on
Commit
7f5fe96
1 Parent(s): 958efd4

Update train_model.py

Browse files
Files changed (1) hide show
  1. src/models/train_model.py +5 -1
src/models/train_model.py CHANGED
@@ -13,4 +13,8 @@ def train_model():
13
  model = Summarization()
14
  model.from_pretrained('t5','t5-base')
15
  model.train(train_df=train_df, eval_df=eval_df, batch_size=4, max_epochs=3, use_gpu=True)
16
- model.save_model()
 
 
 
 
 
13
  model = Summarization()
14
  model.from_pretrained('t5','t5-base')
15
  model.train(train_df=train_df, eval_df=eval_df, batch_size=4, max_epochs=3, use_gpu=True)
16
+ model.save_model()
17
+
18
+
19
+ if __name__ == '__main__':
20
+ train_model()