Spaces:
Running on Zero
Running on Zero
| from typing import Dict | |
| from transformers import AutoTokenizer | |
| MULTI_LANG_TAGS = [ | |
| "<|Chinese|>", "<|English|>", "<|Cantonese|>", | |
| "<|Japanese|>", "<|Korean|>", "<|Spanish|>", | |
| "<|French|>", "<|Russian|>", "<|Arabic|>", | |
| "<|Turkish|>", "<|Indonesian|>", "<|Portuguese|>", | |
| "<|Italian|>", "<|Dutch|>", "<|Vietnamese|>", | |
| "<|German|>", "<|Ukrainian|>", "<|Thai|>", | |
| "<|Polish|>", "<|Romanian|>", "<|Greek|>", | |
| "<|Czech|>", "<|Finnish|>", "<|Hindi|>", | |
| ] | |
| MULTI_DIALECT_TAGS = [ | |
| "<|ZH_Anhui|>", "<|ZH_Fujian|>", "<|ZH_Gansu|>", | |
| "<|ZH_Guizhou|>", "<|ZH_Hebei|>", "<|ZH_Henan|>", | |
| "<|ZH_Hubei|>", "<|ZH_Hunan|>", "<|ZH_Jiangxi|>", | |
| "<|ZH_Liaoning|>", "<|ZH_Minnan|>", "<|ZH_Ningxia|>", | |
| "<|ZH_Shaanxi|>", "<|ZH_Shandong|>", "<|ZH_Shanghai|>", | |
| "<|ZH_Shanxi|>", "<|ZH_Sichuan|>", "<|ZH_Tianjin|>", | |
| "<|ZH_Wenzhou|>", "<|ZH_Wu|>", "<|ZH_Yunnan|>", | |
| ] | |
| def load_text_tokenizer(tokenizer_dir:str): | |
| text_tokenizer = AutoTokenizer.from_pretrained( | |
| pretrained_model_name_or_path=tokenizer_dir | |
| ) | |
| vocab: Dict = text_tokenizer.get_vocab() | |
| # Add special tokens | |
| special_tokens = [ | |
| "<|sosp|>", | |
| "<|eosp|>", | |
| "<|empty|>", | |
| "<|Human|>", | |
| "<|SpeechLM|>", | |
| "<|sostm|>", | |
| "<|eostm|>", | |
| "<|sot|>", | |
| "<|eot|>", | |
| "<|TEXT_ONLY|>", | |
| "<|AUDIO_ONLY|>", | |
| "<|ASR|>", | |
| "<|TTS|>", | |
| "<|INTERLEAVE|>", | |
| "<|UNDERSTANDING|>", | |
| # Placeholders | |
| "<|placeholder_001|>", | |
| "<|placeholder_002|>", | |
| "<|placeholder_003|>", | |
| "<|placeholder_004|>", | |
| "<|placeholder_005|>", | |
| "<|placeholder_006|>", | |
| "<|placeholder_007|>", | |
| "<|placeholder_008|>", | |
| "<|placeholder_009|>", | |
| "<|placeholder_010|>", | |
| "<|placeholder_011|>", | |
| "<|placeholder_012|>", | |
| "<|placeholder_013|>", | |
| "<|placeholder_014|>", | |
| "<|placeholder_015|>", | |
| "<|placeholder_016|>", | |
| "<|placeholder_017|>", | |
| "<|placeholder_018|>", | |
| "<|placeholder_019|>", | |
| "<|placeholder_020|>", | |
| "<|placeholder_021|>", | |
| "<|placeholder_022|>", | |
| "<|placeholder_023|>", | |
| "<|placeholder_024|>", | |
| "<|placeholder_025|>", | |
| "<|placeholder_026|>", | |
| "<|placeholder_027|>", | |
| "<|placeholder_028|>", | |
| "<|placeholder_029|>", | |
| "<|placeholder_030|>", | |
| "<|placeholder_031|>", | |
| "<|placeholder_032|>", | |
| "<|placeholder_033|>", | |
| "<|placeholder_034|>", | |
| "<|placeholder_035|>", | |
| "<|placeholder_036|>", | |
| "<|placeholder_037|>", | |
| "<|placeholder_038|>", | |
| "<|placeholder_039|>", | |
| "<|placeholder_040|>", | |
| "<|placeholder_041|>", | |
| "<|placeholder_042|>", | |
| "<|placeholder_043|>", | |
| "<|placeholder_044|>", | |
| "<|placeholder_045|>", | |
| "<|placeholder_046|>", | |
| "<|placeholder_047|>", | |
| "<|placeholder_048|>", | |
| "<|placeholder_049|>", | |
| "<|placeholder_050|>", | |
| "<|placeholder_051|>", | |
| "<|placeholder_052|>", | |
| "<|placeholder_053|>", | |
| "<|placeholder_054|>", | |
| "<|placeholder_055|>", | |
| "<|placeholder_056|>", | |
| "<|placeholder_057|>", | |
| "<|placeholder_058|>", | |
| "<|placeholder_059|>", | |
| "<|placeholder_060|>", | |
| "<|placeholder_061|>", | |
| "<|placeholder_062|>", | |
| "<|placeholder_063|>", | |
| "<|placeholder_064|>", | |
| "<|placeholder_065|>", | |
| "<|placeholder_066|>", | |
| "<|placeholder_067|>", | |
| "<|placeholder_068|>", | |
| "<|placeholder_069|>", | |
| "<|placeholder_070|>", | |
| "<|placeholder_071|>", | |
| "<|placeholder_072|>", | |
| "<|placeholder_073|>", | |
| "<|placeholder_074|>", | |
| "<|placeholder_075|>", | |
| "<|placeholder_076|>", | |
| "<|placeholder_077|>", | |
| "<|placeholder_078|>", | |
| "<|placeholder_079|>", | |
| "<|placeholder_080|>", | |
| "<|placeholder_081|>", | |
| "<|placeholder_082|>", | |
| "<|placeholder_083|>", | |
| "<|placeholder_084|>", | |
| "<|placeholder_085|>", | |
| "<|placeholder_086|>", | |
| "<|placeholder_087|>", | |
| "<|placeholder_088|>", | |
| "<|placeholder_089|>", | |
| "<|placeholder_090|>", | |
| "<|placeholder_091|>", | |
| "<|placeholder_092|>", | |
| "<|placeholder_093|>", | |
| "<|placeholder_094|>", | |
| "<|placeholder_095|>", | |
| "<|placeholder_096|>", | |
| "<|placeholder_097|>", | |
| "<|placeholder_098|>", | |
| "<|placeholder_099|>", | |
| "<|placeholder_100|>", | |
| "<|placeholder_101|>", | |
| "<|placeholder_102|>", | |
| "<|placeholder_103|>", | |
| "<|placeholder_104|>", | |
| "<|placeholder_105|>", | |
| "<|placeholder_106|>", | |
| "<|placeholder_107|>", | |
| "<|placeholder_108|>", | |
| "<|placeholder_109|>", | |
| "<|placeholder_110|>", | |
| "<|placeholder_111|>", | |
| "<|placeholder_112|>", | |
| "<|placeholder_113|>", | |
| "<|placeholder_114|>", | |
| "<|placeholder_115|>", | |
| "<|placeholder_116|>", | |
| "<|placeholder_117|>", | |
| "<|placeholder_118|>", | |
| "<|placeholder_119|>", | |
| "<|placeholder_120|>", | |
| "<|placeholder_121|>", | |
| "<|placeholder_122|>", | |
| "<|placeholder_123|>", | |
| "<|placeholder_124|>", | |
| "<|placeholder_125|>", | |
| "<|placeholder_126|>", | |
| "<|placeholder_127|>", | |
| "<|placeholder_128|>", | |
| "<|placeholder_129|>", | |
| "<|placeholder_130|>", | |
| "<|placeholder_131|>", | |
| "<|placeholder_132|>", | |
| "<|placeholder_133|>", | |
| "<|placeholder_134|>", | |
| "<|placeholder_135|>", | |
| "<|placeholder_136|>", | |
| "<|placeholder_137|>", | |
| "<|placeholder_138|>", | |
| "<|placeholder_139|>", | |
| "<|placeholder_140|>", | |
| "<|placeholder_141|>", | |
| "<|placeholder_142|>", | |
| "<|placeholder_143|>", | |
| "<|placeholder_144|>", | |
| "<|placeholder_145|>", | |
| "<|placeholder_146|>", | |
| "<|placeholder_147|>", | |
| "<|placeholder_148|>", | |
| "<|placeholder_149|>", | |
| "<|placeholder_150|>", | |
| "<|placeholder_151|>", | |
| "<|placeholder_152|>", | |
| "<|placeholder_153|>", | |
| "<|placeholder_154|>", | |
| "<|placeholder_155|>", | |
| "<|placeholder_156|>", | |
| "<|placeholder_157|>", | |
| "<|placeholder_158|>", | |
| "<|placeholder_159|>", | |
| "<|placeholder_160|>", | |
| "<|placeholder_161|>", | |
| "<|placeholder_162|>", | |
| "<|placeholder_163|>", | |
| "<|placeholder_164|>", | |
| "<|placeholder_165|>", | |
| "<|placeholder_166|>", | |
| "<|placeholder_167|>", | |
| "<|placeholder_168|>", | |
| "<|placeholder_169|>", | |
| "<|placeholder_170|>", | |
| "<|placeholder_171|>", | |
| "<|placeholder_172|>", | |
| "<|placeholder_173|>", | |
| "<|placeholder_174|>", | |
| "<|placeholder_175|>", | |
| "<|placeholder_176|>", | |
| "<|placeholder_177|>", | |
| "<|placeholder_178|>", | |
| "<|placeholder_179|>", | |
| "<|placeholder_180|>", | |
| "<|placeholder_181|>", | |
| "<|placeholder_182|>", | |
| "<|placeholder_183|>", | |
| "<|placeholder_184|>", | |
| "<|placeholder_185|>", | |
| "<|placeholder_186|>", | |
| "<|placeholder_187|>", | |
| "<|placeholder_188|>", | |
| "<|placeholder_189|>", | |
| "<|placeholder_190|>", | |
| "<|placeholder_191|>", | |
| "<|placeholder_192|>", | |
| # Multi-lang tag | |
| *MULTI_LANG_TAGS, | |
| # Multi-dialect tag | |
| *MULTI_DIALECT_TAGS, | |
| # For SpeechEdit CoT | |
| "<|edit|>", "<|frame_patch|>", "<|end_edit|>", | |
| ] | |
| for token in special_tokens: | |
| if token not in vocab: | |
| text_tokenizer.add_tokens([token], special_tokens=True) | |
| return text_tokenizer | |