Gagan Bhatia commited on
Commit
958efd4
1 Parent(s): 62ce5f5

Update train_model.py

Browse files
Files changed (1) hide show
  1. src/models/train_model.py +1 -1
src/models/train_model.py CHANGED
@@ -11,6 +11,6 @@ def train_model():
11
  eval_df = pd.read_csv('../../data/processed/validation.csv')
12
 
13
  model = Summarization()
14
- model.from_pretrained('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()
 
11
  eval_df = pd.read_csv('../../data/processed/validation.csv')
12
 
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()