def get_config(dataset: str, dataset_id: int, model_name: str, epochs: int,
docs: Dict) -> Tuple[TrainConfig, EvalConfig]:
if dataset == 'cookpad':
from src.data.cookpad.preprocessors import ConcatDataProcessor
data_processor = ConcatDataProcessor(docs)
train_config, eval_config = {
'ebr': cookpad.ebr_config,
'naive': cookpad.naive_config,
'nrmf_simple_query': cookpad.nrmf_simple_query_config,
'nrmf_simple_all': cookpad.nrmf_simple_all_config,
'nrmf_simple_query_with_1st':
cookpad.nrmf_simple_query_with_1st_config,
'fwfm_query': cookpad.fwfm_query_config,
'fwfm_all': cookpad.fwfm_all_config,
'fwfm_selected': cookpad.fwfm_selected_config,
'fwfm_all_without_1st': cookpad.fwfm_all_without_1st_config,
}[model_name](f'{dataset}.{dataset_id}', epochs, data_processor)
else:
from src.data.msmarco.preprocessors import ConcatDataProcessor
data_processor = ConcatDataProcessor(docs)
train_config, eval_config = {
'ebr': msmarco.ebr_config,
'naive': msmarco.naive_config,
'nrmf_simple_query': msmarco.nrmf_simple_query_config,
'nrmf_simple_all': msmarco.nrmf_simple_all_config,
'fwfm_query': msmarco.fwfm_query_config,
'fwfm_all': msmarco.fwfm_all_config,
}[model_name](f'{dataset}.{dataset_id}', epochs, data_processor)
return train_config, eval_config
def nrmf_simple_all_config(dataset_id: str, epochs: int,
data_processor) -> Tuple[TrainConfig, EvalConfig]:
if not data_processor:
data_processor = ConcatDataProcessor()
train_config = TrainConfig(
dataset_id=dataset_id,
data_processor=data_processor,
model=nrmf.NRMFSimpleAll,
epochs=epochs,
verbose=2,
)
eval_config = EvalConfig(
dataset_id=dataset_id,
model_name='nrmf_simple_all',
verbose=0,
)
return train_config, eval_config
def ebr_config(
dataset_id: str, epochs: int,
data_processor: DataProcessor) -> Tuple[TrainConfig, EvalConfig]:
if not data_processor:
data_processor = ConcatDataProcessor()
train_config = TrainConfig(
dataset_id=dataset_id,
data_processor=data_processor,
model=representation.EBR,
epochs=epochs,
verbose=2,
)
eval_config = EvalConfig(
dataset_id=dataset_id,
model_name='ebr',
verbose=0,
)
return train_config, eval_config
def naive_config(
dataset_id: str, epochs: int,
data_processor: DataProcessor) -> Tuple[TrainConfig, EvalConfig]:
if not data_processor:
data_processor = ConcatDataProcessor()
train_config = TrainConfig(
dataset_id=dataset_id,
data_processor=data_processor,
model=naive.Naive,
epochs=epochs,
verbose=2,
)
eval_config = EvalConfig(
dataset_id=dataset_id,
model_name='naive',
verbose=0,
)
return train_config, eval_config
def fwfm_all_without_1st_config(
dataset_id: str, epochs: int,
data_processor: DataProcessor) -> Tuple[TrainConfig, EvalConfig]:
if not data_processor:
data_processor = ConcatDataProcessor()
train_config = TrainConfig(
dataset_id=dataset_id,
data_processor=data_processor,
model=fm.FwFMAllWithout1st,
epochs=epochs,
verbose=2,
)
eval_config = EvalConfig(
dataset_id=dataset_id,
model_name='fwfm_all_without_1st',
verbose=0,
)
return train_config, eval_config
def fwfm_selected_config(
dataset_id: str, epochs: int,
data_processor: DataProcessor) -> Tuple[TrainConfig, EvalConfig]:
if not data_processor:
data_processor = ConcatDataProcessor()
train_config = TrainConfig(
dataset_id=dataset_id,
data_processor=data_processor,
model=fm.FwFMSelected,
epochs=epochs,
verbose=2,
)
eval_config = EvalConfig(
dataset_id=dataset_id,
model_name='fwfm_selected',
verbose=0,
)
return train_config, eval_config