Spaces:
Runtime error
Runtime error
Gagan Bhatia
commited on
Commit
•
958efd4
1
Parent(s):
62ce5f5
Update train_model.py
Browse files
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()
|