|
import base64 |
|
import os |
|
from functools import lru_cache |
|
from typing import Optional |
|
import torch |
|
from transformers import AutoTokenizer |
|
from whisper.tokenizer import Tokenizer |
|
|
|
import tiktoken |
|
|
|
LANGUAGES = { |
|
"en": "english", |
|
"zh": "chinese", |
|
"de": "german", |
|
"es": "spanish", |
|
"ru": "russian", |
|
"ko": "korean", |
|
"fr": "french", |
|
"ja": "japanese", |
|
"pt": "portuguese", |
|
"tr": "turkish", |
|
"pl": "polish", |
|
"ca": "catalan", |
|
"nl": "dutch", |
|
"ar": "arabic", |
|
"sv": "swedish", |
|
"it": "italian", |
|
"id": "indonesian", |
|
"hi": "hindi", |
|
"fi": "finnish", |
|
"vi": "vietnamese", |
|
"he": "hebrew", |
|
"uk": "ukrainian", |
|
"el": "greek", |
|
"ms": "malay", |
|
"cs": "czech", |
|
"ro": "romanian", |
|
"da": "danish", |
|
"hu": "hungarian", |
|
"ta": "tamil", |
|
"no": "norwegian", |
|
"th": "thai", |
|
"ur": "urdu", |
|
"hr": "croatian", |
|
"bg": "bulgarian", |
|
"lt": "lithuanian", |
|
"la": "latin", |
|
"mi": "maori", |
|
"ml": "malayalam", |
|
"cy": "welsh", |
|
"sk": "slovak", |
|
"te": "telugu", |
|
"fa": "persian", |
|
"lv": "latvian", |
|
"bn": "bengali", |
|
"sr": "serbian", |
|
"az": "azerbaijani", |
|
"sl": "slovenian", |
|
"kn": "kannada", |
|
"et": "estonian", |
|
"mk": "macedonian", |
|
"br": "breton", |
|
"eu": "basque", |
|
"is": "icelandic", |
|
"hy": "armenian", |
|
"ne": "nepali", |
|
"mn": "mongolian", |
|
"bs": "bosnian", |
|
"kk": "kazakh", |
|
"sq": "albanian", |
|
"sw": "swahili", |
|
"gl": "galician", |
|
"mr": "marathi", |
|
"pa": "punjabi", |
|
"si": "sinhala", |
|
"km": "khmer", |
|
"sn": "shona", |
|
"yo": "yoruba", |
|
"so": "somali", |
|
"af": "afrikaans", |
|
"oc": "occitan", |
|
"ka": "georgian", |
|
"be": "belarusian", |
|
"tg": "tajik", |
|
"sd": "sindhi", |
|
"gu": "gujarati", |
|
"am": "amharic", |
|
"yi": "yiddish", |
|
"lo": "lao", |
|
"uz": "uzbek", |
|
"fo": "faroese", |
|
"ht": "haitian creole", |
|
"ps": "pashto", |
|
"tk": "turkmen", |
|
"nn": "nynorsk", |
|
"mt": "maltese", |
|
"sa": "sanskrit", |
|
"lb": "luxembourgish", |
|
"my": "myanmar", |
|
"bo": "tibetan", |
|
"tl": "tagalog", |
|
"mg": "malagasy", |
|
"as": "assamese", |
|
"tt": "tatar", |
|
"haw": "hawaiian", |
|
"ln": "lingala", |
|
"ha": "hausa", |
|
"ba": "bashkir", |
|
"jw": "javanese", |
|
"su": "sundanese", |
|
"yue": "cantonese", |
|
"minnan": "minnan", |
|
"wuyu": "wuyu", |
|
"dialect": "dialect", |
|
"zh/en": "zh/en", |
|
"en/zh": "en/zh", |
|
} |
|
|
|
|
|
TO_LANGUAGE_CODE = { |
|
**{language: code for code, language in LANGUAGES.items()}, |
|
"burmese": "my", |
|
"valencian": "ca", |
|
"flemish": "nl", |
|
"haitian": "ht", |
|
"letzeburgesch": "lb", |
|
"pushto": "ps", |
|
"panjabi": "pa", |
|
"moldavian": "ro", |
|
"moldovan": "ro", |
|
"sinhalese": "si", |
|
"castilian": "es", |
|
"mandarin": "zh", |
|
} |
|
|
|
AUDIO_EVENT = { |
|
"ASR": "ASR", |
|
"AED": "AED", |
|
"SER": "SER", |
|
"Speech": "Speech", |
|
"/Speech": "/Speech", |
|
"BGM": "BGM", |
|
"/BGM": "/BGM", |
|
"Laughter": "Laughter", |
|
"/Laughter": "/Laughter", |
|
"Applause": "Applause", |
|
"/Applause": "/Applause", |
|
} |
|
|
|
EMOTION = { |
|
"HAPPY": "HAPPY", |
|
"SAD": "SAD", |
|
"ANGRY": "ANGRY", |
|
"NEUTRAL": "NEUTRAL", |
|
} |
|
|
|
TTS_Vocal_Token = { |
|
"TTS/B": "TTS/B", |
|
"TTS/O": "TTS/O", |
|
"TTS/Q": "TTS/Q", |
|
"TTS/A": "TTS/A", |
|
"TTS/CO": "TTS/CO", |
|
"TTS/CL": "TTS/CL", |
|
"TTS/H": "TTS/H", |
|
**{f"TTS/SP{i:02d}": f"TTS/SP{i:02d}" for i in range(1, 14)} |
|
} |
|
|
|
|
|
@lru_cache(maxsize=None) |
|
def get_encoding(name: str = "gpt2", num_languages: int = 99): |
|
vocab_path = os.path.join(os.path.dirname(__file__), "assets", f"{name}.tiktoken") |
|
ranks = { |
|
base64.b64decode(token): int(rank) |
|
for token, rank in (line.split() for line in open(vocab_path) if line) |
|
} |
|
n_vocab = len(ranks) |
|
special_tokens = {} |
|
|
|
specials = [ |
|
"<|endoftext|>", |
|
"<|startoftranscript|>", |
|
*[f"<|{lang}|>" for lang in list(LANGUAGES.keys())[:num_languages]], |
|
*[f"<|{audio_event}|>" for audio_event in list(AUDIO_EVENT.keys())], |
|
*[f"<|{emotion}|>" for emotion in list(EMOTION.keys())], |
|
"<|translate|>", |
|
"<|transcribe|>", |
|
"<|startoflm|>", |
|
"<|startofprev|>", |
|
"<|nospeech|>", |
|
"<|notimestamps|>", |
|
*[f"<|SPECIAL_TOKEN_{i}|>" for i in range(1, 31)], |
|
*[f"<|{tts}|>" for tts in list(TTS_Vocal_Token.keys())], |
|
*[f"<|{i * 0.02:.2f}|>" for i in range(1501)], |
|
] |
|
|
|
for token in specials: |
|
special_tokens[token] = n_vocab |
|
n_vocab += 1 |
|
|
|
return tiktoken.Encoding( |
|
name=os.path.basename(vocab_path), |
|
explicit_n_vocab=n_vocab, |
|
pat_str=r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""", |
|
mergeable_ranks=ranks, |
|
special_tokens=special_tokens, |
|
) |
|
|
|
|
|
@lru_cache(maxsize=None) |
|
def get_tokenizer( |
|
multilingual: bool, |
|
*, |
|
num_languages: int = 99, |
|
language: Optional[str] = None, |
|
task: Optional[str] = None, |
|
) -> Tokenizer: |
|
if language is not None: |
|
language = language.lower() |
|
if language not in LANGUAGES: |
|
if language in TO_LANGUAGE_CODE: |
|
language = TO_LANGUAGE_CODE[language] |
|
else: |
|
raise ValueError(f"Unsupported language: {language}") |
|
|
|
if multilingual: |
|
encoding_name = "multilingual_zh_ja_yue_char_del" |
|
language = language or "en" |
|
task = task or "transcribe" |
|
else: |
|
encoding_name = "gpt2" |
|
language = None |
|
task = None |
|
|
|
encoding = get_encoding(name=encoding_name, num_languages=num_languages) |
|
|
|
return Tokenizer( |
|
encoding=encoding, num_languages=num_languages, language=language, task=task |
|
) |
|
|
|
|
|
class QwenTokenizer(): |
|
def __init__(self, token_path, skip_special_tokens=True): |
|
super().__init__() |
|
|
|
special_tokens = { |
|
'eos_token': '<|endoftext|>', |
|
'pad_token': '<|endoftext|>', |
|
'additional_special_tokens': [ |
|
'<|im_start|>', '<|im_end|>', '<|endofprompt|>', |
|
'[breath]', '<strong>', '</strong>', '[noise]', |
|
'[laughter]', '[cough]', '[clucking]', '[accent]', |
|
'[quick_breath]', |
|
"<laughter>", "</laughter>", |
|
"[hissing]", "[sigh]", "[vocalized-noise]", |
|
"[lipsmack]", "[mn]" |
|
] |
|
} |
|
self.special_tokens = special_tokens |
|
self.tokenizer = AutoTokenizer.from_pretrained(token_path) |
|
self.tokenizer.add_special_tokens(special_tokens) |
|
self.skip_special_tokens = skip_special_tokens |
|
|
|
def encode(self, text, **kwargs): |
|
tokens = self.tokenizer([text], return_tensors="pt") |
|
tokens = tokens["input_ids"][0].cpu().tolist() |
|
return tokens |
|
|
|
def decode(self, tokens): |
|
tokens = torch.tensor(tokens, dtype=torch.int64) |
|
text = self.tokenizer.batch_decode([tokens], skip_special_tokens=self.skip_special_tokens)[0] |
|
return text |
|
|
|
|
|
@lru_cache(maxsize=None) |
|
def get_qwen_tokenizer( |
|
token_path: str, |
|
skip_special_tokens: bool |
|
) -> QwenTokenizer: |
|
return QwenTokenizer(token_path=token_path, skip_special_tokens=skip_special_tokens) |
|
|
