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