# -*- coding: utf-8 -*-
try:
    from .kronos import KronosTokenizer, Kronos, KronosPredictor
    model_dict = {
        'kronos_tokenizer': KronosTokenizer,
        'kronos': Kronos,
        'kronos_predictor': KronosPredictor
    }
    def get_model_class(model_name):
        if model_name in model_dict:
            return model_dict[model_name]
        else:
            raise NotImplementedError(f"Model {model_name} not found")
except ImportError:
    model_dict = {}
    def get_model_class(model_name):
        raise NotImplementedError(f"Model {model_name} not available (torch not installed)")