FilesExpand file tree

master

/

__init__.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

17 lines (12 loc) · 412 Bytes

master

/

__init__.py

Copy path

Top

File metadata and controls

• Code
• Blame

17 lines (12 loc) · 412 Bytes

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

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:

print(f"Model {model_name} not found in model_dict")

raise NotImplementedError