T5-Summarization / src /models /train_model.py
Gagan Bhatia
Update train_model.py
4c60028
raw
history blame
572 Bytes
import yaml
from src.models.model import Summarization
import pandas as pd
def train_model():
"""
Train the model
"""
with open("params.yml") as f:
params = yaml.safe_load(f)
# Load the data
train_df = pd.read_csv('data/processed/train.csv')
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()
if __name__ == '__main__':
train_model()