164 lines
4.6 KiB
Python
164 lines
4.6 KiB
Python
# /// script
|
|
# requires-python = ">=3.12"
|
|
# dependencies = [
|
|
# "pypinyin",
|
|
# ]
|
|
# [tool.uv]
|
|
# exclude-newer = "2025-05-06T00:00:00Z"
|
|
# ///
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
from typing import List, Union
|
|
from pypinyin import pinyin
|
|
from pypinyin.contrib.tone_convert import to_initials, to_finals_tone
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--text",
|
|
type=str,
|
|
required=False,
|
|
default="my-keywords.txt",
|
|
help="""Path to the input texts.
|
|
|
|
Each line in the texts contains the original phrase, it might also contain some
|
|
extra items, for example, the boosting score (startting with :), the triggering
|
|
threshold (startting with #, only used in keyword spotting task) and the original
|
|
phrase (startting with @). Note: extra items will be kept in the output.
|
|
|
|
example input 1 (tokens_type = ppinyin):
|
|
|
|
小爱同学 :2.0 #0.6 @小爱同学
|
|
你好问问 :3.5 @你好问问
|
|
小艺小艺 #0.6 @小艺小艺
|
|
|
|
example output 1:
|
|
|
|
x iǎo ài t óng x ué :2.0 #0.6 @小爱同学
|
|
n ǐ h ǎo w èn w èn :3.5 @你好问问
|
|
x iǎo y ì x iǎo y ì #0.6 @小艺小艺
|
|
|
|
example input 2 (tokens_type = bpe):
|
|
|
|
HELLO WORLD :1.5 #0.4
|
|
HI GOOGLE :2.0 #0.8
|
|
HEY SIRI #0.35
|
|
|
|
example output 2:
|
|
|
|
▁HE LL O ▁WORLD :1.5 #0.4
|
|
▁HI ▁GO O G LE :2.0 #0.8
|
|
▁HE Y ▁S I RI #0.35
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--tokens",
|
|
type=str,
|
|
required=False,
|
|
default="tokens.txt",
|
|
help="The path to tokens.txt.",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--output",
|
|
type=str,
|
|
required=False,
|
|
default="keywords.txt",
|
|
help="Path where the encoded tokens will be written to.",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def text2token(
|
|
texts: List[str],
|
|
tokens: str,
|
|
output_ids: bool = False,
|
|
):
|
|
"""将文本转换为 token 列表"""
|
|
assert Path(tokens).is_file(), f"File not exists, {tokens}"
|
|
tokens_table = {}
|
|
with open(tokens, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
toks = line.strip().split()
|
|
assert len(toks) == 2, len(toks)
|
|
assert toks[0] not in tokens_table, f"Duplicate token: {toks} "
|
|
tokens_table[toks[0]] = int(toks[1])
|
|
|
|
texts_list: List[List[str]] = []
|
|
for txt in texts:
|
|
res = []
|
|
py = [x[0] for x in pinyin(txt)]
|
|
for x in py:
|
|
initial = to_initials(x, strict=False)
|
|
final = to_finals_tone(x, strict=False)
|
|
if initial == "" and final == "":
|
|
res.append(x)
|
|
else:
|
|
if initial != "":
|
|
res.append(initial)
|
|
if final != "":
|
|
res.append(final)
|
|
texts_list.append(res)
|
|
|
|
result: List[List[Union[int, str]]] = []
|
|
for text in texts_list:
|
|
text_list = []
|
|
contain_oov = False
|
|
for txt in text:
|
|
if txt in tokens_table:
|
|
text_list.append(tokens_table[txt] if output_ids else txt)
|
|
else:
|
|
print(
|
|
f"Can't find token {txt} in token table, check your "
|
|
f"tokens.txt see if {txt} in it. skipping text : {text}."
|
|
)
|
|
contain_oov = True
|
|
break
|
|
if contain_oov:
|
|
continue
|
|
else:
|
|
result.append(text_list)
|
|
return result
|
|
|
|
|
|
def main() -> None:
|
|
args = get_args()
|
|
|
|
texts = []
|
|
# extra information like boosting score (start with :), triggering threshold (start with #)
|
|
# original keyword (start with @)
|
|
extra_info = []
|
|
with open(args.text, "r", encoding="utf8") as f:
|
|
for line in f:
|
|
extra = []
|
|
text = []
|
|
toks = line.strip().split()
|
|
if len(toks) == 1:
|
|
text.append(toks[0])
|
|
extra.append('@'+toks[0])
|
|
else:
|
|
for tok in toks:
|
|
if tok[0] == ":" or tok[0] == "#" or tok[0] == "@":
|
|
extra.append(tok)
|
|
else:
|
|
text.append(tok)
|
|
texts.append(" ".join(text))
|
|
extra_info.append(extra)
|
|
encoded_texts = text2token(
|
|
texts,
|
|
tokens=args.tokens,
|
|
)
|
|
with open(args.output, "w", encoding="utf8") as f:
|
|
for i, txt in enumerate(encoded_texts):
|
|
txt += extra_info[i]
|
|
f.write(" ".join(txt) + "\n")
|
|
print(f"✅ 唤醒词已保存到 {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|