-
Notifications
You must be signed in to change notification settings - Fork 31
/
model_map.py
41 lines (40 loc) · 1.71 KB
/
model_map.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
from . import model_utils
from . import tokenizer_utils
MODEL_MAP = {
'albert':
{'resizer': model_utils.AlbertVocabResizer,
'tokenizer_helper': tokenizer_utils.SentencepieceTokenizer,
'structure': model_utils.AlbertStructure},
'bert':
{'resizer': model_utils.BertVocabResizer,
'tokenizer_helper': tokenizer_utils.SubwordTokenizer,
'structure': model_utils.BertStructure},
'electra':
{'resizer': model_utils.ElectraVocabResizer,
'tokenizer_helper': tokenizer_utils.SubwordTokenizer,
'structure': model_utils.ElectraStructure},
'roberta':
{'resizer': model_utils.RobertaVocabResizer,
'tokenizer_helper' : tokenizer_utils.RobertaGPT2Tokenizer,
'structure': model_utils.RobertaStructure},
'xlm-roberta':
{'resizer':model_utils.XLMRobertaVocabResizer,
'tokenizer_helper': tokenizer_utils.XLMRSentencepieceTokenizer,
'structure': model_utils.XLMRobertaStructure},
'xlm':
{'resizer':model_utils.XLMVocabResizer,
'tokenizer_helper':tokenizer_utils.XLMTokenizer,
'structure':model_utils.XLMStructure},
'bart':
{'resizer' : model_utils.BartVocabResizer,
'tokenizer_helper' : tokenizer_utils.RobertaGPT2Tokenizer,
'structure': model_utils.BartStructure},
't5':
{'resizer' : model_utils.T5VocabResizer,
'tokenizer_helper' : tokenizer_utils.T5SentencepieceTokenizer,
'structure' : model_utils.T5Structure},
'mt5':
{'resizer' : model_utils.MT5VocabResizer,
'tokenizer_helper' : tokenizer_utils.MT5SentencepieceTokenizer,
'structure' : model_utils.MT5Structure},
}