File size: 2,460 Bytes
90f0b29 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import logging
import re
import sentencepiece as spm
from fireredasr.data.token_dict import TokenDict
class ChineseCharEnglishSpmTokenizer:
"""
- One Chinese char is a token.
- Split English word into SPM and one piece is a token.
- Ignore ' ' between Chinese char
- Replace ' ' between English word with "▁" by spm_model
- Need to put SPM piece into dict file
- If not set spm_model, will use English char and <space>
"""
SPM_SPACE = "▁"
def __init__(self, dict_path, spm_model, unk="<unk>", space="<space>"):
self.dict = TokenDict(dict_path, unk=unk)
self.space = space
if spm_model:
self.sp = spm.SentencePieceProcessor()
self.sp.Load(spm_model)
else:
self.sp = None
print("[WRAN] Not set spm_model, will use English char")
print("[WARN] Please check how to deal with ' '(space)")
if self.space not in self.dict:
print("Please add <space> to your dict, or it will be <unk>")
def tokenize(self, text, replace_punc=True):
#if text == "":
# logging.info(f"empty text")
text = text.upper()
tokens = []
if replace_punc:
text = re.sub("[，。？！,\.?!]", " ", text)
pattern = re.compile(r'([\u3400-\u4dbf\u4e00-\u9fff])')
parts = pattern.split(text.strip())
parts = [p for p in parts if len(p.strip()) > 0]
for part in parts:
if pattern.fullmatch(part) is not None:
tokens.append(part)
else:
if self.sp:
for piece in self.sp.EncodeAsPieces(part.strip()):
tokens.append(piece)
else:
for char in part.strip():
tokens.append(char if char != " " else self.space)
tokens_id = []
for token in tokens:
tokens_id.append(self.dict.get(token, self.dict.unk))
return tokens, tokens_id
def detokenize(self, inputs, join_symbol="", replace_spm_space=True):
"""inputs is ids or tokens, do not need self.sp"""
if len(inputs) > 0 and type(inputs[0]) == int:
tokens = [self.dict[id] for id in inputs]
else:
tokens = inputs
s = f"{join_symbol}".join(tokens)
if replace_spm_space:
s = s.replace(self.SPM_SPACE, ' ').strip()
return s
|