Files
ai-toolkit/toolkit/job.py
2023-08-01 13:49:54 -06:00

25 lines
668 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)
if job == 'mod':
from jobs import ModJob
return ModJob(config)
# elif job == 'train':
# from jobs import TrainJob
# return TrainJob(config)
else:
raise ValueError(f'Unknown job type {job}')