372 lines
14 KiB
Python
372 lines
14 KiB
Python
import numpy as np
|
|
from funasr.register import tables
|
|
import logging
|
|
import random
|
|
import re
|
|
|
|
|
|
@tables.register("prompt_classes", "MultiContextPrompt")
|
|
class MultiContextPrompt:
|
|
CONTEXT_TEMPLATES = {
|
|
'en': {
|
|
'header': "Please combine the context information provided below to complete the speech transcription task more accurately. If there is no relevant information, we will leave it blank.\n",
|
|
'fields': {
|
|
'hist_context': "Historical transcription: {hist_context}\n",
|
|
'one_pass_result': "One-pass result: {one_pass_result}\n",
|
|
'hotwords': "Hotword list: {hotwords}\n"
|
|
}
|
|
},
|
|
'zh': {
|
|
'header': "请结合下面提供的上下文信息,更加准确地完成语音转写任务。如果没有相关信息,我们会留空。\n",
|
|
'fields': {
|
|
'hist_context': "历史转写结果:{hist_context}\n",
|
|
'one_pass_result': "一遍解码结果:{one_pass_result}\n",
|
|
'hotwords': "热词列表:{hotwords}\n"
|
|
}
|
|
}
|
|
}
|
|
|
|
def __init__(self,
|
|
use_hist=True,
|
|
use_one_pass_result=True,
|
|
use_hotwords=True,
|
|
use_asr_hotwords=True,
|
|
use_multi_lingual_prompt=True,
|
|
**kwargs):
|
|
"""Initialize MultiContextPrompt.
|
|
|
|
Args:
|
|
use_hist: TODO.
|
|
use_one_pass_result: TODO.
|
|
use_hotwords: TODO.
|
|
use_asr_hotwords: TODO.
|
|
use_multi_lingual_prompt: TODO.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
self.use_hist = use_hist
|
|
self.use_one_pass_result = use_one_pass_result
|
|
self.use_hotwords = use_hotwords
|
|
self.use_asr_hotwords = use_asr_hotwords
|
|
self.use_multi_lingual_prompt = use_multi_lingual_prompt
|
|
self.kwargs = kwargs
|
|
|
|
chinese_hotwords_list = kwargs.get("chinese_hotwords_list", "")
|
|
english_hotwords_list = kwargs.get("english_hotwords_list", "")
|
|
if chinese_hotwords_list:
|
|
self.chinese_hotwords_list, self.chinese_hotwords_num = self.get_hotwords_list(chinese_hotwords_list)
|
|
else:
|
|
self.chinese_hotwords_list = None
|
|
self.chinese_hotwords_num = 0
|
|
logging.info(f"chinese_hotwords_num: {self.chinese_hotwords_num}")
|
|
|
|
if english_hotwords_list:
|
|
self.english_hotwords_list, self.english_hotwords_num = self.get_hotwords_list(english_hotwords_list)
|
|
else:
|
|
self.english_hotwords_list = None
|
|
self.english_hotwords_num = 0
|
|
logging.info(f"english_hotwords_num: {self.english_hotwords_num}")
|
|
|
|
self.max_neg_hotwords_num = kwargs.get("max_neg_hotwords_num", 900)
|
|
self.min_neg_hotwords_num = kwargs.get("min_neg_hotwords_num", 0)
|
|
|
|
def get_hotwords_list(self, hotwords_file):
|
|
"""Get hotwords list.
|
|
|
|
Args:
|
|
hotwords_file: TODO.
|
|
"""
|
|
with open(hotwords_file, "r") as f:
|
|
hotwords_list = f.read().strip().split("\n")
|
|
return hotwords_list, len(hotwords_list)
|
|
|
|
def detect_language(self, text):
|
|
"""Detect language.
|
|
|
|
Args:
|
|
text: Text tensor or string input.
|
|
"""
|
|
if isinstance(text, list):
|
|
text = " ".join(text)
|
|
|
|
chinese_pattern = re.compile(
|
|
"["
|
|
"\u4e00-\u9fff" # CJK Unified Ideographs
|
|
"]+"
|
|
)
|
|
|
|
english_pattern = re.compile(r'[A-Za-z]+')
|
|
|
|
chinese_matches = chinese_pattern.findall(text)
|
|
english_matches = english_pattern.findall(text)
|
|
|
|
chinese_length = sum(len(match) for match in chinese_matches)
|
|
english_length = sum(len(match) for match in english_matches)
|
|
|
|
total_length = len(text)
|
|
|
|
if total_length == 0:
|
|
return 'zh'
|
|
|
|
if (chinese_length > english_length) and (chinese_length / total_length > 0.3):
|
|
return 'zh'
|
|
else:
|
|
return 'en'
|
|
|
|
def hotwords_sampling(self, hotwords):
|
|
|
|
# hotwords_list = hotwords.split(", ")
|
|
"""Hotwords sampling.
|
|
|
|
Args:
|
|
hotwords: TODO.
|
|
"""
|
|
hotwords_list = hotwords
|
|
selected_hotwords = []
|
|
if self.max_neg_hotwords_num > -1:
|
|
max_neg_hotwords_num = min(self.max_neg_hotwords_num, len(hotwords_list))
|
|
else:
|
|
max_neg_hotwords_num = len(hotwords_list)
|
|
|
|
if self.min_neg_hotwords_num < max_neg_hotwords_num:
|
|
selected_hotwords_num = np.random.randint(self.min_neg_hotwords_num, max_neg_hotwords_num + 1)
|
|
else:
|
|
selected_hotwords_num = max_neg_hotwords_num
|
|
if selected_hotwords_num > 0:
|
|
selected_hotwords = np.random.choice(hotwords_list, selected_hotwords_num, replace=False).tolist()
|
|
|
|
return selected_hotwords, selected_hotwords_num
|
|
|
|
def get_prompt(self, item, language):
|
|
"""Get prompt.
|
|
|
|
Args:
|
|
item: TODO.
|
|
language: Language identifier.
|
|
"""
|
|
template = self.CONTEXT_TEMPLATES[language]
|
|
|
|
prompt = template['header']
|
|
|
|
context_lines = []
|
|
|
|
if self.use_hist and item.get("hist_context"):
|
|
context_lines.append(template['fields']['hist_context'].format(hist_context=item["hist_context"]))
|
|
|
|
if self.use_one_pass_result and item.get("one_pass_result"):
|
|
context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=item["one_pass_result"]))
|
|
|
|
hotwords = None
|
|
if self.use_hotwords and item.get("hotwords"):
|
|
hotwords = item["hotwords"]
|
|
if self.use_asr_hotwords and item.get("asr_hotwords"):
|
|
hotwords = item["asr_hotwords"]
|
|
if hotwords is not None and hotwords != "":
|
|
language = self.detect_language(hotwords)
|
|
if language == 'en':
|
|
neg_hotwords = self.english_hotwords_list
|
|
else:
|
|
neg_hotwords = self.chinese_hotwords_list
|
|
if neg_hotwords is not None:
|
|
selected_neg_hotwords, selected_neg_hotwords_num = self.hotwords_sampling(neg_hotwords)
|
|
else:
|
|
selected_neg_hotwords = []
|
|
|
|
if not isinstance(hotwords, list):
|
|
pos_hotwords = hotwords.split(", ")
|
|
else:
|
|
pos_hotwords = hotwords
|
|
hotwords = pos_hotwords + selected_neg_hotwords
|
|
random.shuffle(hotwords)
|
|
hotwords = ", ".join(hotwords)
|
|
context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
|
|
|
|
if context_lines:
|
|
prompt += ''.join(context_lines)
|
|
else:
|
|
prompt += "\n\n\n"
|
|
|
|
return prompt
|
|
|
|
def get_inference_prompt(self, item, language="zh"):
|
|
"""Get inference prompt.
|
|
|
|
Args:
|
|
item: TODO.
|
|
language: Language identifier.
|
|
"""
|
|
template = self.CONTEXT_TEMPLATES[language]
|
|
|
|
prompt = template['header']
|
|
|
|
context_lines = []
|
|
|
|
if self.use_hist and item.get("hist_context"):
|
|
context_lines.append(template['fields']['hist_context'].format(hist_context=item["hist_context"]))
|
|
|
|
if self.use_one_pass_result and item.get("one_pass_result"):
|
|
context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=item["one_pass_result"]))
|
|
|
|
hotwords = None
|
|
if self.use_hotwords and item.get("hotwords"):
|
|
hotwords = item["hotwords"]
|
|
if self.use_asr_hotwords and item.get("asr_hotwords"):
|
|
hotwords = item["asr_hotwords"]
|
|
if hotwords is not None and hotwords != "":
|
|
print(f"hotwords: {hotwords}")
|
|
language = self.detect_language(hotwords)
|
|
if language == 'en':
|
|
neg_hotwords = self.english_hotwords_list
|
|
else:
|
|
neg_hotwords = self.chinese_hotwords_list
|
|
if neg_hotwords is not None:
|
|
selected_neg_hotwords, selected_neg_hotwords_num = self.hotwords_sampling(neg_hotwords)
|
|
else:
|
|
selected_neg_hotwords = []
|
|
|
|
if not isinstance(hotwords, list):
|
|
pos_hotwords = hotwords.split(", ")
|
|
else:
|
|
pos_hotwords = hotwords
|
|
hotwords = pos_hotwords + selected_neg_hotwords
|
|
print(f"selected_neg_hotwords_num: {selected_neg_hotwords_num}")
|
|
random.shuffle(hotwords)
|
|
hotwords = ", ".join(hotwords)
|
|
context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
|
|
|
|
if context_lines:
|
|
prompt += ''.join(context_lines)
|
|
else:
|
|
prompt += "\n\n\n"
|
|
|
|
return prompt
|
|
|
|
|
|
@tables.register("prompt_classes", "MultiContextPromptNew")
|
|
class MultiContextPromptNew:
|
|
CONTEXT_TEMPLATES = {
|
|
'en': {
|
|
'header': "Please combine the context information to complete the speech transcription task more accurately. If there is no relevant information, we will leave it blank.\n\n",
|
|
'context_header': "**Context:**\n",
|
|
'fields': {
|
|
'hist_context': "Historical transcription: {hist_context}\n",
|
|
'one_pass_result': "One-pass result: {one_pass_result}\n",
|
|
'hotwords': "Hotword list: {hotwords}\n"
|
|
}
|
|
},
|
|
'zh': {
|
|
'header': "请结合上下文信息,更加准确地完成语音转写任务。如果没有相关信息,我们会留空。\n\n",
|
|
'context_header': "**上下文信息:**\n",
|
|
'fields': {
|
|
'hist_context': "历史转写结果:{hist_context}\n",
|
|
'one_pass_result': "一遍解码结果:{one_pass_result}\n",
|
|
'hotwords': "热词列表:{hotwords}\n"
|
|
}
|
|
}
|
|
}
|
|
|
|
def __init__(self,
|
|
use_hist=True,
|
|
use_one_pass_result=True,
|
|
use_hotwords=True,
|
|
use_multi_lingual_prompt=True,
|
|
**kwargs):
|
|
"""Initialize MultiContextPromptNew.
|
|
|
|
Args:
|
|
use_hist: TODO.
|
|
use_one_pass_result: TODO.
|
|
use_hotwords: TODO.
|
|
use_multi_lingual_prompt: TODO.
|
|
**kwargs: Additional keyword arguments.
|
|
"""
|
|
self.use_hist = use_hist
|
|
self.use_one_pass_result = use_one_pass_result
|
|
self.use_hotwords = use_hotwords
|
|
self.use_multi_lingual_prompt = use_multi_lingual_prompt
|
|
|
|
self.use_full_hotwords_ratio = kwargs.get("use_full_hotwords_ratio", 0.2)
|
|
self.max_hotwords_num = kwargs.get("max_hotwords_num", -1)
|
|
self.min_hotwords_num = kwargs.get("min_hotwords_num", 15)
|
|
|
|
def hotwords_sampling(self, hotwords):
|
|
|
|
"""Hotwords sampling.
|
|
|
|
Args:
|
|
hotwords: TODO.
|
|
"""
|
|
hotwords_list = hotwords.split(", ")
|
|
if self.max_hotwords_num > 0:
|
|
max_hotwords_num = min(self.max_hotwords_num, len(hotwords_list))
|
|
else:
|
|
max_hotwords_num = len(hotwords_list)
|
|
|
|
if self.min_hotwords_num < max_hotwords_num:
|
|
selected_hotwords_num = np.random.randint(self.min_hotwords_num, max_hotwords_num + 1)
|
|
else:
|
|
selected_hotwords_num = max_hotwords_num
|
|
|
|
selected_hotwords = np.random.choice(hotwords_list, selected_hotwords_num, replace=False)
|
|
hotwords_list = ", ".join(selected_hotwords)
|
|
|
|
return hotwords_list, selected_hotwords_num
|
|
|
|
def get_prompt(self, item, language):
|
|
"""Get prompt.
|
|
|
|
Args:
|
|
item: TODO.
|
|
language: Language identifier.
|
|
"""
|
|
template = self.CONTEXT_TEMPLATES[language]
|
|
|
|
prompt = template['header']
|
|
|
|
context_lines = []
|
|
|
|
if self.use_hist and item.get("hist_context"):
|
|
context_lines.append(template['fields']['hist_context'].format(hist_context=item["hist_context"]))
|
|
|
|
if self.use_one_pass_result and item.get("one_pass_result"):
|
|
context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=item["one_pass_result"]))
|
|
|
|
if self.use_hotwords and item.get("hotwords"):
|
|
hotwords = item["hotwords"]
|
|
if np.random.rand() < self.use_full_hotwords_ratio:
|
|
hotwords = hotwords
|
|
else:
|
|
hotwords, selected_hotwords_num = self.hotwords_sampling(hotwords)
|
|
context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
|
|
|
|
if context_lines:
|
|
prompt += template['context_header'] + ''.join(context_lines)
|
|
|
|
return prompt
|
|
|
|
def get_inference_prompt(self, hist_context="", one_pass_result="", hotwords=""):
|
|
"""Get inference prompt.
|
|
|
|
Args:
|
|
hist_context: TODO.
|
|
one_pass_result: TODO.
|
|
hotwords: TODO.
|
|
"""
|
|
language = 'zh' if self.use_multi_lingual_prompt and np.random.rand() < 0.5 else 'en'
|
|
template = self.CONTEXT_TEMPLATES[language]
|
|
|
|
prompt = template['header']
|
|
|
|
context_lines = []
|
|
|
|
if hist_context:
|
|
context_lines.append(template['fields']['hist_context'].format(hist_context=hist_context))
|
|
if one_pass_result:
|
|
context_lines.append(template['fields']['one_pass_result'].format(one_pass_result=one_pass_result))
|
|
if hotwords:
|
|
context_lines.append(template['fields']['hotwords'].format(hotwords=hotwords))
|
|
|
|
if context_lines:
|
|
prompt += template['context_header'] + ''.join(context_lines)
|
|
|
|
return prompt
|