Files
ai-toolkit/toolkit/job.py
2023-07-31 11:56:57 -06:00

22 lines
585 B
Python

from toolkit.config import get_config
def get_job(config_path, name=None):
config = get_config(config_path, name)
if not config['job']:
raise ValueError('config file is invalid. Missing "job" key')
job = config['job']
if job == 'extract':
from jobs import ExtractJob
return ExtractJob(config)
if job == 'train':
from jobs import TrainJob
return TrainJob(config)
# elif job == 'train':
# from jobs import TrainJob
# return TrainJob(config)
else:
raise ValueError(f'Unknown job type {job}')