107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
import json
|
|
from pathlib import Path
|
|
import re
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
from handle.text_handle import process_text
|
|
from handle.text_normalizer import UYGHUR_LETTERS
|
|
|
|
class ASRTokenizer:
|
|
def __init__(self, vocab_path: Path) -> None:
|
|
self.vocabs: list[str] = []
|
|
self.vocabs.extend([
|
|
"<BLANK>",
|
|
"<UNK>",
|
|
"<SPACE>"
|
|
])
|
|
for i in range(len(self.vocabs), 64):
|
|
self.vocabs.append(f"REVERSED_{i}")
|
|
|
|
with open(vocab_path, 'r', encoding='utf-8') as file:
|
|
data: list[str] = json.load(file)
|
|
|
|
for char in data:
|
|
if char not in [' ', '\t', '\n'] and char not in self.vocabs:
|
|
self.vocabs.append(char)
|
|
|
|
self.vocab_to_id: dict[str, int] = {vocab:i for i, vocab in enumerate(self.vocabs)}
|
|
self.id_to_vocab: dict[int, str] = {i:vocab for i, vocab in enumerate(self.vocabs)}
|
|
|
|
self._max_token_len = max((len(v) for v in self.vocabs if not v.startswith("<")), default=0)
|
|
|
|
def vocab_size(self):
|
|
return len(self.vocabs)
|
|
|
|
def _split_long_syllable(self, syll: str) -> list[str]:
|
|
"""将不在词表中的长音节拆分为词表内的子词序列（最大正向匹配）"""
|
|
pieces = []
|
|
start = 0
|
|
n = len(syll)
|
|
while start < n:
|
|
# 从最大可能长度开始尝试匹配
|
|
for length in range(min(self._max_token_len, n - start), 0, -1):
|
|
sub = syll[start:start + length]
|
|
if sub in self.vocab_to_id:
|
|
pieces.append(sub)
|
|
start += length
|
|
break
|
|
else:
|
|
# 理论上不会发生，因为单字母一定在词表
|
|
pieces.append("<UNK>")
|
|
start += 1
|
|
return pieces
|
|
|
|
def encode(self, text: str) -> list[int]:
|
|
tokens = process_text(text=text)
|
|
result: list[int] = []
|
|
|
|
for token in tokens:
|
|
if token == ' ':
|
|
result.append(self.vocab_to_id['<SPACE>'])
|
|
else:
|
|
token_id = self.vocab_to_id.get(token)
|
|
if token_id is not None:
|
|
result.append(token_id)
|
|
elif token in UYGHUR_LETTERS:
|
|
result.append(self.vocab_to_id.get(token, self.vocab_to_id['<UNK>']))
|
|
else:
|
|
sub_pieces = self._split_long_syllable(token)
|
|
for piece in sub_pieces:
|
|
result.append(self.vocab_to_id.get(piece, self.vocab_to_id['<UNK>']))
|
|
|
|
return result
|
|
|
|
def decode(self, ids: list[int] | Tensor, remove_blank: bool = True, remove_repeat: bool = True) -> str:
|
|
if isinstance(ids, Tensor):
|
|
ids = ids.tolist()
|
|
|
|
result = []
|
|
prev_id = None
|
|
|
|
for token_id in ids:
|
|
# 跳过blank token
|
|
if remove_blank and token_id == self.vocab_to_id['<BLANK>']:
|
|
prev_id = None
|
|
continue
|
|
|
|
# CTC解码：跳过连续重复的字符
|
|
if remove_repeat and token_id == prev_id:
|
|
continue
|
|
|
|
char = self.id_to_vocab.get(token_id, "<UNK>")
|
|
|
|
if char == '<SPACE>':
|
|
char = ' '
|
|
|
|
result.append(char)
|
|
prev_id = token_id
|
|
|
|
return ''.join(result)
|
|
|
|
def get_special_token_id(self, token: str) -> int:
|
|
return self.vocab_to_id.get(token, self.vocab_to_id['<UNK>'])
|
|
|
|
def ctc_greedy_decode(self, log_probs: Tensor) -> str:
|
|
""" CTC贪心解码 (log_probs: [seq_len, vocab_size] 的log概率) """
|
|
return self.decode(torch.argmax(log_probs, dim=-1)) |