Spaces:
Running
Running
Commit
·
5a751c2
1
Parent(s):
dc54ed7
优化LLM消息格式化和响应处理逻辑,增强错误处理和调试信息,改进说话人识别器的JSON输出提取逻辑。
Browse files
src/podcast_transcribe/llm/llm_base.py
CHANGED
@@ -22,78 +22,178 @@ class BaseChatCompletion(ABC):
|
|
22 |
pass
|
23 |
|
24 |
def _format_messages_for_gemma(self, messages: List[Dict[str, str]]) -> str:
|
25 |
-
"""
|
26 |
-
为Gemma格式化消息。
|
27 |
-
Gemma期望特定的格式,通常类似于:
|
28 |
-
<start_of_turn>user
|
29 |
-
{user_message}<end_of_turn>
|
30 |
-
<start_of_turn>model
|
31 |
-
{assistant_message}<end_of_turn>
|
32 |
-
...
|
33 |
-
<start_of_turn>user
|
34 |
-
{current_user_message}<end_of_turn>
|
35 |
-
<start_of_turn>model
|
36 |
-
"""
|
37 |
-
# 尝试使用分词器的聊天模板(如果可用)
|
38 |
try:
|
39 |
-
#
|
40 |
-
|
41 |
-
|
42 |
-
|
43 |
-
|
44 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
45 |
)
|
46 |
-
|
47 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
48 |
prompt_parts = []
|
49 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
50 |
role = message.get("role")
|
51 |
-
content = message.get("content")
|
|
|
|
|
|
|
|
|
52 |
if role == "user":
|
53 |
prompt_parts.append(f"<start_of_turn>user\n{content}<end_of_turn>")
|
54 |
elif role == "assistant":
|
55 |
prompt_parts.append(f"<start_of_turn>model\n{content}<end_of_turn>")
|
56 |
-
|
57 |
-
|
58 |
-
# 我们会在这里前置它,尽管其有效性取决于特定的Gemma微调。
|
59 |
-
# 一种常见的模式是在开头放置系统指令,不使用特殊标记。
|
60 |
-
# 然而,为了保持结构化,我们将尝试一种通用方法。
|
61 |
-
# 如果分词器在其模板中有特定的方式来处理系统提示,
|
62 |
-
# 那么`apply_chat_template`将是首选。
|
63 |
-
# 由于我们处于回退状态,这是一个最佳猜测。
|
64 |
-
# 一些模型期望系统提示在轮次结构之外,或者在最开始。
|
65 |
-
# 为了在回退中简化,我们只做前置处理。
|
66 |
-
# 如果`apply_chat_template`不可用,更健壮的解决方案是检查模型的特定聊天模板。
|
67 |
-
prompt_parts.insert(0, f"<start_of_turn>system\n{content}<end_of_turn>")
|
68 |
-
|
69 |
-
# 添加提示,让模型开始生成
|
70 |
prompt_parts.append("<start_of_turn>model")
|
71 |
-
|
|
|
|
|
|
|
72 |
|
73 |
def _post_process_response(self, response_text: str, prompt_str: str) -> str:
|
74 |
"""
|
75 |
后处理生成的响应文本,清理提示和特殊标记
|
76 |
"""
|
77 |
-
|
78 |
-
|
79 |
-
|
80 |
-
#
|
81 |
-
|
82 |
-
# 如果模型输出包含提示,然后是新的响应:
|
83 |
if response_text.startswith(prompt_str):
|
84 |
assistant_message_content = response_text[len(prompt_str):].strip()
|
|
|
85 |
else:
|
86 |
-
#
|
87 |
-
|
88 |
-
|
89 |
-
|
90 |
-
parts = response_text.split("<start_of_turn>model")
|
91 |
-
if len(parts) > 1:
|
92 |
assistant_message_content = parts[-1].strip()
|
93 |
-
|
94 |
-
|
95 |
-
|
96 |
assistant_message_content = response_text.strip()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
97 |
|
98 |
return assistant_message_content
|
99 |
|
@@ -223,19 +323,26 @@ class TransformersBaseChatCompletion(BaseChatCompletion):
|
|
223 |
print("警告: MPS 设备不支持 device_map,将手动管理设备")
|
224 |
else:
|
225 |
model_kwargs["device_map"] = self.device_map
|
|
|
226 |
|
227 |
# 加载模型
|
|
|
228 |
self.model = AutoModelForCausalLM.from_pretrained(
|
229 |
self.model_name,
|
230 |
**model_kwargs
|
231 |
)
|
232 |
|
233 |
# MPS 或手动设备管理
|
234 |
-
if self.device_map is None:
|
235 |
print(f"手动移动模型到设备: {self.device}")
|
236 |
self.model = self.model.to(self.device)
|
237 |
|
|
|
|
|
|
|
238 |
print(f"模型 {self.model_name} 加载成功")
|
|
|
|
|
239 |
|
240 |
def _load_model_and_tokenizer(self):
|
241 |
"""加载模型和分词器"""
|
@@ -269,26 +376,102 @@ class TransformersBaseChatCompletion(BaseChatCompletion):
|
|
269 |
inputs = self.tokenizer.encode(prompt_str, return_tensors="pt")
|
270 |
|
271 |
# 移动输入到正确的设备
|
272 |
-
if self.device_map is None or (self.device and self.device.type == "mps"):
|
273 |
inputs = inputs.to(self.device)
|
274 |
|
275 |
-
#
|
276 |
generation_config = {
|
277 |
"max_new_tokens": max_tokens,
|
278 |
-
"temperature": temperature,
|
279 |
-
"top_p": top_p,
|
280 |
-
"do_sample": True if temperature > 0 else False,
|
281 |
"pad_token_id": self.tokenizer.pad_token_id,
|
282 |
"eos_token_id": self.tokenizer.eos_token_id,
|
283 |
-
"
|
284 |
-
"no_repeat_ngram_size": kwargs.get("no_repeat_ngram_size", 3),
|
285 |
}
|
286 |
|
287 |
-
#
|
288 |
-
if temperature
|
289 |
-
|
290 |
-
|
291 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
292 |
|
293 |
try:
|
294 |
# 生成响应
|
@@ -302,10 +485,41 @@ class TransformersBaseChatCompletion(BaseChatCompletion):
|
|
302 |
generated_tokens = outputs[0][len(inputs[0]):]
|
303 |
generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
|
304 |
|
|
|
305 |
return generated_text
|
306 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
307 |
except Exception as e:
|
308 |
print(f"生成响应时出错: {e}")
|
|
|
|
|
309 |
raise
|
310 |
|
311 |
def get_model_info(self) -> Dict[str, Union[str, bool, int]]:
|
|
|
22 |
pass
|
23 |
|
24 |
def _format_messages_for_gemma(self, messages: List[Dict[str, str]]) -> str:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
25 |
try:
|
26 |
+
# 确保消息格式正确
|
27 |
+
formatted_messages = []
|
28 |
+
for msg in messages:
|
29 |
+
if msg.get("role") and msg.get("content"):
|
30 |
+
formatted_messages.append({
|
31 |
+
"role": msg["role"],
|
32 |
+
"content": msg["content"]
|
33 |
+
})
|
34 |
+
|
35 |
+
# 使用官方聊天模板
|
36 |
+
prompt_str = self.tokenizer.apply_chat_template(
|
37 |
+
formatted_messages,
|
38 |
+
tokenize=False,
|
39 |
+
add_generation_prompt=True
|
40 |
)
|
41 |
+
|
42 |
+
# 调试信息
|
43 |
+
print(f"使用官方聊天模板格式化成功,长度: {len(prompt_str)}")
|
44 |
+
return prompt_str
|
45 |
+
|
46 |
+
except Exception as e:
|
47 |
+
print(f"官方聊天模板失败: {e},使用手动格式化")
|
48 |
+
|
49 |
+
# 手动格式化 - 改进版本
|
50 |
prompt_parts = []
|
51 |
+
|
52 |
+
# 处理系统消息 - Gemma 3 的正确处理方式
|
53 |
+
system_messages = [msg for msg in messages if msg.get("role") == "system"]
|
54 |
+
other_messages = [msg for msg in messages if msg.get("role") != "system"]
|
55 |
+
|
56 |
+
# 对于 Gemma,系统消息通常需要特殊处理
|
57 |
+
if system_messages:
|
58 |
+
# 将系统消息作为第一个用户消息的前缀
|
59 |
+
system_content = "\n".join([msg["content"] for msg in system_messages])
|
60 |
+
if other_messages and other_messages[0].get("role") == "user":
|
61 |
+
# 将系统提示合并到第一个用户消息中
|
62 |
+
first_user_msg = other_messages[0]
|
63 |
+
combined_content = f"{system_content}\n\n{first_user_msg['content']}"
|
64 |
+
other_messages[0] = {"role": "user", "content": combined_content}
|
65 |
+
else:
|
66 |
+
# 如果没有用户消息,创建一个包含系统提示的用户消息
|
67 |
+
other_messages.insert(0, {"role": "user", "content": system_content})
|
68 |
+
|
69 |
+
# 格式化其他消息
|
70 |
+
for message in other_messages:
|
71 |
role = message.get("role")
|
72 |
+
content = message.get("content", "").strip()
|
73 |
+
|
74 |
+
if not content:
|
75 |
+
continue
|
76 |
+
|
77 |
if role == "user":
|
78 |
prompt_parts.append(f"<start_of_turn>user\n{content}<end_of_turn>")
|
79 |
elif role == "assistant":
|
80 |
prompt_parts.append(f"<start_of_turn>model\n{content}<end_of_turn>")
|
81 |
+
|
82 |
+
# 添加生成提示
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
83 |
prompt_parts.append("<start_of_turn>model")
|
84 |
+
|
85 |
+
formatted_prompt = "\n".join(prompt_parts)
|
86 |
+
print(f"手动格式化完成,长度: {len(formatted_prompt)}")
|
87 |
+
return formatted_prompt
|
88 |
|
89 |
def _post_process_response(self, response_text: str, prompt_str: str) -> str:
|
90 |
"""
|
91 |
后处理生成的响应文本,清理提示和特殊标记
|
92 |
"""
|
93 |
+
print(f"原始响应长度: {len(response_text)}")
|
94 |
+
print(f"原始响应前100字符: {response_text[:100]}")
|
95 |
+
|
96 |
+
# 如果模型输出包含提示,则移除提示部分
|
|
|
|
|
97 |
if response_text.startswith(prompt_str):
|
98 |
assistant_message_content = response_text[len(prompt_str):].strip()
|
99 |
+
print("检测到响应包含提示,已移除提示部分")
|
100 |
else:
|
101 |
+
# 尝试找到最后一个 "<start_of_turn>model" 标记
|
102 |
+
model_start_marker = "<start_of_turn>model"
|
103 |
+
if model_start_marker in response_text:
|
104 |
+
parts = response_text.split(model_start_marker)
|
|
|
|
|
105 |
assistant_message_content = parts[-1].strip()
|
106 |
+
print("通过 <start_of_turn>model 标记分割响应")
|
107 |
+
else:
|
108 |
+
# 如果没有找到标记,使用整个响应
|
109 |
assistant_message_content = response_text.strip()
|
110 |
+
print("未找到特殊标记,使用完整响应")
|
111 |
+
|
112 |
+
# 清理结束标记
|
113 |
+
end_markers = ["<end_of_turn>", "<|endoftext|>", "</s>"]
|
114 |
+
for marker in end_markers:
|
115 |
+
if marker in assistant_message_content:
|
116 |
+
assistant_message_content = assistant_message_content.split(marker)[0].strip()
|
117 |
+
print(f"移除结束标记: {marker}")
|
118 |
+
|
119 |
+
# 特殊处理:如果响应看起来像 JSON,尝试提取 JSON 部分
|
120 |
+
if "{" in assistant_message_content and "}" in assistant_message_content:
|
121 |
+
# 尝试提取 JSON 对象
|
122 |
+
first_brace = assistant_message_content.find("{")
|
123 |
+
last_brace = assistant_message_content.rfind("}")
|
124 |
+
if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
|
125 |
+
potential_json = assistant_message_content[first_brace:last_brace + 1]
|
126 |
+
# 验证是否为有效 JSON
|
127 |
+
try:
|
128 |
+
import json
|
129 |
+
json.loads(potential_json)
|
130 |
+
assistant_message_content = potential_json
|
131 |
+
print("提取并验证了 JSON 响应")
|
132 |
+
except json.JSONDecodeError:
|
133 |
+
print("JSON 验证失败,保持原始响应")
|
134 |
+
# 如果JSON验证失败,尝试清理常见的非JSON内容
|
135 |
+
lines = assistant_message_content.split('\n')
|
136 |
+
cleaned_lines = []
|
137 |
+
json_started = False
|
138 |
+
|
139 |
+
for line in lines:
|
140 |
+
line = line.strip()
|
141 |
+
# 跳过明显的解释性文本
|
142 |
+
if any(phrase in line.lower() for phrase in [
|
143 |
+
'here is', 'here\'s', 'based on', 'analysis', 'looking at',
|
144 |
+
'the json', 'response:', 'result:', 'output:', 'answer:',
|
145 |
+
'i can see', 'it appears', 'according to'
|
146 |
+
]):
|
147 |
+
continue
|
148 |
+
|
149 |
+
# 如果遇到JSON开始,标记开始
|
150 |
+
if '{' in line:
|
151 |
+
json_started = True
|
152 |
+
|
153 |
+
# 如果已经开始JSON,保留所有内容
|
154 |
+
if json_started:
|
155 |
+
cleaned_lines.append(line)
|
156 |
+
|
157 |
+
# 如果遇到JSON结束,停止
|
158 |
+
if '}' in line and json_started:
|
159 |
+
break
|
160 |
+
|
161 |
+
if cleaned_lines:
|
162 |
+
assistant_message_content = '\n'.join(cleaned_lines)
|
163 |
+
print("清理了非JSON解释性内容")
|
164 |
+
|
165 |
+
# 最终清理:移除常见的解释性前缀和后缀
|
166 |
+
prefixes_to_remove = [
|
167 |
+
"Here is the JSON:", "Here's the JSON:", "The JSON response is:",
|
168 |
+
"Based on the analysis:", "Looking at the information:",
|
169 |
+
"Here is my analysis:", "The result is:", "My response:",
|
170 |
+
"Output:", "Answer:", "Result:"
|
171 |
+
]
|
172 |
+
|
173 |
+
for prefix in prefixes_to_remove:
|
174 |
+
if assistant_message_content.lower().startswith(prefix.lower()):
|
175 |
+
assistant_message_content = assistant_message_content[len(prefix):].strip()
|
176 |
+
print(f"移除前缀: {prefix}")
|
177 |
+
break
|
178 |
+
|
179 |
+
# 移除常见的后缀解释
|
180 |
+
suffixes_to_remove = [
|
181 |
+
"This JSON object maps each speaker ID to their identified name or role.",
|
182 |
+
"Each speaker has been identified based on the provided information.",
|
183 |
+
"The identification is based on the dialogue samples and metadata."
|
184 |
+
]
|
185 |
+
|
186 |
+
for suffix in suffixes_to_remove:
|
187 |
+
if assistant_message_content.lower().endswith(suffix.lower()):
|
188 |
+
assistant_message_content = assistant_message_content[:-len(suffix)].strip()
|
189 |
+
print(f"移除后缀: {suffix}")
|
190 |
+
break
|
191 |
+
|
192 |
+
# 最终清理
|
193 |
+
assistant_message_content = assistant_message_content.strip()
|
194 |
+
|
195 |
+
print(f"处理后响应长度: {len(assistant_message_content)}")
|
196 |
+
print(f"处理后响应前100字符: {assistant_message_content[:100]}")
|
197 |
|
198 |
return assistant_message_content
|
199 |
|
|
|
323 |
print("警告: MPS 设备不支持 device_map,将手动管理设备")
|
324 |
else:
|
325 |
model_kwargs["device_map"] = self.device_map
|
326 |
+
print(f"使用设备映射: {self.device_map}")
|
327 |
|
328 |
# 加载模型
|
329 |
+
print("开始加载模型...")
|
330 |
self.model = AutoModelForCausalLM.from_pretrained(
|
331 |
self.model_name,
|
332 |
**model_kwargs
|
333 |
)
|
334 |
|
335 |
# MPS 或手动设备管理
|
336 |
+
if self.device_map is None and self.device is not None:
|
337 |
print(f"手动移动模型到设备: {self.device}")
|
338 |
self.model = self.model.to(self.device)
|
339 |
|
340 |
+
# 设置模型为评估模式
|
341 |
+
self.model.eval()
|
342 |
+
|
343 |
print(f"模型 {self.model_name} 加载成功")
|
344 |
+
print(f"模型数据类型: {self.model.dtype}")
|
345 |
+
print(f"模型设备: {next(self.model.parameters()).device}")
|
346 |
|
347 |
def _load_model_and_tokenizer(self):
|
348 |
"""加载模型和分词器"""
|
|
|
376 |
inputs = self.tokenizer.encode(prompt_str, return_tensors="pt")
|
377 |
|
378 |
# 移动输入到正确的设备
|
379 |
+
if self.device_map is None or (self.device and hasattr(self.device, 'type') and self.device.type == "mps"):
|
380 |
inputs = inputs.to(self.device)
|
381 |
|
382 |
+
# 优化的生成参数配置
|
383 |
generation_config = {
|
384 |
"max_new_tokens": max_tokens,
|
|
|
|
|
|
|
385 |
"pad_token_id": self.tokenizer.pad_token_id,
|
386 |
"eos_token_id": self.tokenizer.eos_token_id,
|
387 |
+
"use_cache": True, # 启用 KV 缓存以提高速度
|
|
|
388 |
}
|
389 |
|
390 |
+
# 温度和采样配置 - 修复 CUDA 采样错误
|
391 |
+
if temperature > 0:
|
392 |
+
# 确保温度值在合理范围内
|
393 |
+
temperature = max(0.01, min(temperature, 2.0))
|
394 |
+
top_p = max(0.01, min(top_p, 1.0))
|
395 |
+
|
396 |
+
generation_config.update({
|
397 |
+
"do_sample": True,
|
398 |
+
"temperature": temperature,
|
399 |
+
"top_p": top_p,
|
400 |
+
"top_k": kwargs.get("top_k", 10), # 降低top_k以提高确定性
|
401 |
+
})
|
402 |
+
else:
|
403 |
+
# 贪婪解码 - 完全确定性
|
404 |
+
generation_config.update({
|
405 |
+
"do_sample": False,
|
406 |
+
"temperature": None,
|
407 |
+
"top_p": None,
|
408 |
+
"top_k": None,
|
409 |
+
})
|
410 |
+
|
411 |
+
# 如果明确指定do_sample=False,强制使用贪婪解码
|
412 |
+
if kwargs.get("do_sample") is False:
|
413 |
+
generation_config.update({
|
414 |
+
"do_sample": False,
|
415 |
+
"temperature": None,
|
416 |
+
"top_p": None,
|
417 |
+
"top_k": None,
|
418 |
+
})
|
419 |
+
print("强制使用贪婪解码模式")
|
420 |
+
|
421 |
+
# 重复惩罚配置 - 针对结构化输出优化
|
422 |
+
repetition_penalty = kwargs.get("repetition_penalty", 1.0) # 默认不使用重复惩罚
|
423 |
+
if repetition_penalty != 1.0:
|
424 |
+
repetition_penalty = max(1.0, min(repetition_penalty, 1.3)) # 限制在更小范围内
|
425 |
+
generation_config["repetition_penalty"] = repetition_penalty
|
426 |
+
|
427 |
+
# 移除no_repeat_ngram_size以避免干扰JSON格式
|
428 |
+
# generation_config["no_repeat_ngram_size"] = kwargs.get("no_repeat_ngram_size", 2)
|
429 |
+
|
430 |
+
# 长度惩罚(可选)- 针对结构化输出调整
|
431 |
+
if kwargs.get("length_penalty") and kwargs["length_penalty"] != 1.0:
|
432 |
+
length_penalty = max(0.9, min(kwargs["length_penalty"], 1.1)) # 更保守的长度惩罚
|
433 |
+
generation_config["length_penalty"] = length_penalty
|
434 |
+
|
435 |
+
# 针对结构化输出的特殊配置
|
436 |
+
if max_tokens <= 256: # 如果是短输出任务(如JSON),使用更确定性的配置
|
437 |
+
generation_config.update({
|
438 |
+
"early_stopping": True,
|
439 |
+
"num_beams": 1, # 使用贪婪搜索
|
440 |
+
})
|
441 |
+
# 如果允许采样,使用较低的温度
|
442 |
+
if generation_config.get("do_sample", False):
|
443 |
+
generation_config["temperature"] = min(generation_config.get("temperature", 0.1), 0.3)
|
444 |
+
generation_config["top_p"] = min(generation_config.get("top_p", 0.3), 0.5)
|
445 |
+
print("检测到短输出任务,使用优化的生成配置")
|
446 |
+
|
447 |
+
# 处理stop tokens
|
448 |
+
stop_strings = kwargs.get("stop", [])
|
449 |
+
if stop_strings:
|
450 |
+
# 将stop字符串转换为token IDs
|
451 |
+
stop_token_ids = []
|
452 |
+
for stop_str in stop_strings:
|
453 |
+
try:
|
454 |
+
# 编码stop字符串为token IDs
|
455 |
+
stop_tokens = self.tokenizer.encode(stop_str, add_special_tokens=False)
|
456 |
+
stop_token_ids.extend(stop_tokens)
|
457 |
+
except Exception as e:
|
458 |
+
print(f"无法编码stop字符串 '{stop_str}': {e}")
|
459 |
+
|
460 |
+
if stop_token_ids:
|
461 |
+
# 去重并添加到eos_token_id列表中
|
462 |
+
existing_eos = generation_config.get("eos_token_id", self.tokenizer.eos_token_id)
|
463 |
+
if isinstance(existing_eos, int):
|
464 |
+
existing_eos = [existing_eos]
|
465 |
+
elif existing_eos is None:
|
466 |
+
existing_eos = []
|
467 |
+
|
468 |
+
all_stop_tokens = list(set(existing_eos + stop_token_ids))
|
469 |
+
generation_config["eos_token_id"] = all_stop_tokens
|
470 |
+
print(f"添加了 {len(stop_token_ids)} 个stop token IDs")
|
471 |
+
|
472 |
+
# 调试信息
|
473 |
+
print(f"生成配置: temperature={temperature}, max_tokens={max_tokens}, top_p={top_p}")
|
474 |
+
print(f"输入长度: {len(inputs[0])} tokens")
|
475 |
|
476 |
try:
|
477 |
# 生成响应
|
|
|
485 |
generated_tokens = outputs[0][len(inputs[0]):]
|
486 |
generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
|
487 |
|
488 |
+
print(f"生成完成,输出长度: {len(generated_tokens)} tokens")
|
489 |
return generated_text
|
490 |
|
491 |
+
except RuntimeError as e:
|
492 |
+
if "CUDA error" in str(e):
|
493 |
+
print(f"CUDA 错误,尝试使用 CPU 进行推理: {e}")
|
494 |
+
# 尝试移动到 CPU 并重试
|
495 |
+
try:
|
496 |
+
inputs_cpu = inputs.cpu()
|
497 |
+
model_cpu = self.model.cpu()
|
498 |
+
|
499 |
+
with torch.no_grad():
|
500 |
+
outputs = model_cpu.generate(
|
501 |
+
inputs_cpu,
|
502 |
+
**generation_config
|
503 |
+
)
|
504 |
+
|
505 |
+
# 移回原设备
|
506 |
+
self.model = self.model.to(self.device)
|
507 |
+
|
508 |
+
generated_tokens = outputs[0][len(inputs_cpu[0]):]
|
509 |
+
generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
|
510 |
+
|
511 |
+
print(f"CPU 推理完成,输出长度: {len(generated_tokens)} tokens")
|
512 |
+
return generated_text
|
513 |
+
|
514 |
+
except Exception as cpu_e:
|
515 |
+
print(f"CPU 推理也失败: {cpu_e}")
|
516 |
+
raise e
|
517 |
+
else:
|
518 |
+
raise e
|
519 |
except Exception as e:
|
520 |
print(f"生成响应时出错: {e}")
|
521 |
+
import traceback
|
522 |
+
traceback.print_exc()
|
523 |
raise
|
524 |
|
525 |
def get_model_info(self) -> Dict[str, Union[str, bool, int]]:
|
src/podcast_transcribe/summary/speaker_identify.py
CHANGED
@@ -167,67 +167,38 @@ class SpeakerIdentifier:
|
|
167 |
if len(cleaned_episode_shownotes) > max_shownotes_length:
|
168 |
episode_shownotes_for_prompt += "..."
|
169 |
|
170 |
-
system_prompt = """You are
|
171 |
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
189 |
-
|
190 |
-
|
191 |
-
{json.dumps(speaker_info_for_prompt, ensure_ascii=False, indent=2)}
|
192 |
-
```
|
193 |
-
(Analyze dialogue samples and speech statistics to understand speaker roles and identities. DO NOT use speaker IDs to determine roles - SPEAKER_00 is not necessarily the host.)
|
194 |
|
195 |
-
|
196 |
-
|
|
|
197 |
|
198 |
-
|
199 |
-
* A host typically has more frequent, shorter segments, often introduces the show or guests, and may mention the podcast name
|
200 |
-
* In panel discussion formats, there might be multiple hosts or co-hosts of similar speaking patterns
|
201 |
-
* In interview formats, the host typically asks questions while guests give longer answers
|
202 |
-
* Speakers who make introductory statements or welcome listeners are likely hosts
|
203 |
-
* Use dialogue content (not just speaking patterns) to identify names and roles
|
204 |
|
205 |
-
|
206 |
-
|
207 |
-
* The keys of the JSON object MUST be the EXACT "speaker_id"s provided in the input's "Speakers to Identify and Their Information" section (e.g., "SPEAKER_00", "SPEAKER_01").
|
208 |
-
* The values in the JSON object should be the identified person's name or role (string type).
|
209 |
-
* **Prioritize Specific Names/Nicknames**: If there is sufficient information (e.g., guests explicitly listed in Shownotes, or names mentioned in dialogue), please use the identified specific names, such as "John Doe", "AI Assistant", "Dr. Evelyn Reed". Do NOT append roles like "(Host)" or "(Guest)" if a specific name is found.
|
210 |
-
* **Host Identification**:
|
211 |
-
* Hosts may be identified by analyzing speech patterns - they often speak more frequently in shorter segments
|
212 |
-
* Look for introduction patterns in dialogue where speakers welcome listeners or introduce the show
|
213 |
-
* The podcast author (if provided and credible) is often a host but verify through dialogue
|
214 |
-
* There may be multiple hosts (co-hosts) in panel-style podcasts
|
215 |
-
* If a host's name is identified, use the identified name directly (e.g., "Lex Fridman"). Do not append "(Host)".
|
216 |
-
* If the host's name cannot be determined but the role is clearly a host, use "Podcast Host".
|
217 |
-
* **Guest Identification**:
|
218 |
-
* Guests often give longer responses and speak less frequently than hosts
|
219 |
-
* For other non-host speakers, if a specific name is identified, use the identified name directly (e.g., "John Carmack"). Do not append "(Guest)".
|
220 |
-
* If specific names cannot be identified for guests, label them sequentially as "Guest 1", "Guest 2", etc.
|
221 |
-
* **Handling Multiple Hosts/Guests**: If there are multiple hosts or guests and they can be distinguished by name, use their names. If you cannot distinguish specific identities but know there are multiple hosts, use "Host 1", "Host 2", etc. Similarly for guests without specific names, use "Guest 1", "Guest 2".
|
222 |
-
* **Ensure Completeness**: The returned JSON object must include ALL "speaker_id"s listed in the input's "Speakers to Identify and Their Information" section as keys. Each "speaker_id" from the input MUST be a key in your output JSON.
|
223 |
|
224 |
-
|
225 |
-
{{
|
226 |
-
"SPEAKER_00": "Jane Smith",
|
227 |
-
"SPEAKER_01": "Podcast Host",
|
228 |
-
"SPEAKER_02": "Alex Green"
|
229 |
-
}}
|
230 |
-
"""
|
231 |
|
232 |
messages = [
|
233 |
{"role": "system", "content": system_prompt},
|
@@ -281,39 +252,63 @@ Ensure the output is a valid JSON object. For example:
|
|
281 |
messages=messages,
|
282 |
provider=self.llm_provider,
|
283 |
model=self.llm_model_name,
|
284 |
-
temperature=0.
|
285 |
-
max_tokens=
|
286 |
-
|
|
|
|
|
|
|
287 |
)
|
288 |
logger.info(f"LLM调用日志,请求参数:【{messages}】, 响应: 【{response}】")
|
289 |
assistant_response_content = response["choices"][0]["message"]["content"]
|
290 |
|
|
|
291 |
parsed_llm_output = None
|
292 |
-
# 尝试从Markdown代码块中提取JSON
|
293 |
-
json_match = re.search(r'```json\s*(\{.*?\})\s*```', assistant_response_content, re.DOTALL)
|
294 |
-
if json_match:
|
295 |
-
json_str = json_match.group(1)
|
296 |
-
else:
|
297 |
-
# 如果没有markdown块,尝试找到第一个 '{' 到最后一个 '}'
|
298 |
-
first_brace = assistant_response_content.find('{')
|
299 |
-
last_brace = assistant_response_content.rfind('}')
|
300 |
-
if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
|
301 |
-
json_str = assistant_response_content[first_brace : last_brace+1]
|
302 |
-
else: # 如果还是找不到,就认为整个回复都是JSON(可能需要更复杂的清理)
|
303 |
-
json_str = assistant_response_content.strip()
|
304 |
|
|
|
305 |
try:
|
306 |
-
|
307 |
-
|
308 |
-
|
309 |
-
|
310 |
-
|
311 |
-
|
312 |
-
|
313 |
-
|
314 |
-
|
315 |
-
|
316 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
317 |
|
318 |
if parsed_llm_output:
|
319 |
# 直接使用LLM的有效输出,不再依赖预设的角色分配逻辑
|
@@ -346,6 +341,8 @@ Ensure the output is a valid JSON object. For example:
|
|
346 |
|
347 |
if most_active_unknown:
|
348 |
final_map[most_active_unknown] = "Podcast Host"
|
|
|
|
|
349 |
|
350 |
return final_map
|
351 |
|
|
|
167 |
if len(cleaned_episode_shownotes) > max_shownotes_length:
|
168 |
episode_shownotes_for_prompt += "..."
|
169 |
|
170 |
+
system_prompt = """You are a speaker identification expert. Return only a JSON object mapping speaker IDs to names. Start directly with { and end with }. No markdown, no explanations."""
|
171 |
|
172 |
+
# 进一步简化,只保留最关键的信息
|
173 |
+
key_info = []
|
174 |
+
for speaker_id in unique_speaker_ids:
|
175 |
+
samples = dialogue_samples.get(speaker_id, [])
|
176 |
+
stats = speaker_stats.get(speaker_id, {"total_segments": 0, "intro_likely": False})
|
177 |
+
|
178 |
+
# 构建简短描述
|
179 |
+
desc_parts = []
|
180 |
+
if stats["intro_likely"]:
|
181 |
+
desc_parts.append("intro")
|
182 |
+
if stats["total_segments"] > 0:
|
183 |
+
desc_parts.append(f"{stats['total_segments']}segs")
|
184 |
+
if samples:
|
185 |
+
# 只取第一个样本的前50个字符
|
186 |
+
sample_text = samples[0][:50].replace('\n', ' ').strip()
|
187 |
+
if sample_text:
|
188 |
+
desc_parts.append(f'"{sample_text}"')
|
189 |
+
|
190 |
+
key_info.append(f"{speaker_id}: {', '.join(desc_parts)}")
|
|
|
|
|
|
|
191 |
|
192 |
+
user_prompt_template = f"""Podcast: {podcast_title}
|
193 |
+
Host: {podcast_author}
|
194 |
+
Episode: {episode_title}
|
195 |
|
196 |
+
Notes: {episode_shownotes_for_prompt[:300]}
|
|
|
|
|
|
|
|
|
|
|
197 |
|
198 |
+
Speakers:
|
199 |
+
{chr(10).join(key_info)}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
200 |
|
201 |
+
Return JSON like: {{"SPEAKER_00": "Name1", "SPEAKER_01": "Name2"}}"""
|
|
|
|
|
|
|
|
|
|
|
|
|
202 |
|
203 |
messages = [
|
204 |
{"role": "system", "content": system_prompt},
|
|
|
252 |
messages=messages,
|
253 |
provider=self.llm_provider,
|
254 |
model=self.llm_model_name,
|
255 |
+
temperature=0.2, # 稍微提高温度
|
256 |
+
max_tokens=300, # 进一步增加token数
|
257 |
+
top_p=0.5, # 适度提高top_p
|
258 |
+
device=self.device,
|
259 |
+
repetition_penalty=1.0, # 保持不使用重复惩罚
|
260 |
+
do_sample=True # 允许少量采样,不使用stop tokens
|
261 |
)
|
262 |
logger.info(f"LLM调用日志,请求参数:【{messages}】, 响应: 【{response}】")
|
263 |
assistant_response_content = response["choices"][0]["message"]["content"]
|
264 |
|
265 |
+
# 更严格的JSON提取逻辑
|
266 |
parsed_llm_output = None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
267 |
|
268 |
+
# 首先尝试直接解析整个响应(如果它就是JSON)
|
269 |
try:
|
270 |
+
parsed_llm_output = json.loads(assistant_response_content.strip())
|
271 |
+
if isinstance(parsed_llm_output, dict):
|
272 |
+
print("直接解析响应为JSON成功")
|
273 |
+
else:
|
274 |
+
parsed_llm_output = None
|
275 |
+
except json.JSONDecodeError:
|
276 |
+
pass
|
277 |
+
|
278 |
+
# 如果直接解析失败,尝试提取JSON部分
|
279 |
+
if parsed_llm_output is None:
|
280 |
+
# 尝试从Markdown代码块中提取JSON
|
281 |
+
json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', assistant_response_content, re.DOTALL)
|
282 |
+
if json_match:
|
283 |
+
json_str = json_match.group(1)
|
284 |
+
print("从markdown代码块中提取JSON")
|
285 |
+
else:
|
286 |
+
# 如果没有markdown块,尝试找到第一个 '{' 到最后一个 '}'
|
287 |
+
first_brace = assistant_response_content.find('{')
|
288 |
+
last_brace = assistant_response_content.rfind('}')
|
289 |
+
if first_brace != -1 and last_brace != -1 and last_brace > first_brace:
|
290 |
+
json_str = assistant_response_content[first_brace : last_brace+1]
|
291 |
+
print("通过大括号位置提取JSON")
|
292 |
+
else:
|
293 |
+
print("无法找到有效的JSON结构,使用默认映射")
|
294 |
+
return final_map
|
295 |
+
|
296 |
+
try:
|
297 |
+
# 清理JSON字符串
|
298 |
+
json_str = json_str.strip()
|
299 |
+
# 移除可能的换行符和多余空格
|
300 |
+
json_str = re.sub(r'\s+', ' ', json_str)
|
301 |
+
|
302 |
+
parsed_llm_output = json.loads(json_str)
|
303 |
+
if not isinstance(parsed_llm_output, dict):
|
304 |
+
print(f"LLM返回的JSON不是一个字典: {parsed_llm_output}")
|
305 |
+
parsed_llm_output = None
|
306 |
+
else:
|
307 |
+
print("JSON解析成功")
|
308 |
+
except json.JSONDecodeError as e:
|
309 |
+
print(f"LLM返回的JSON解析失败: {e}")
|
310 |
+
print(f"用于解析的字符串: '{json_str[:200]}...'")
|
311 |
+
parsed_llm_output = None
|
312 |
|
313 |
if parsed_llm_output:
|
314 |
# 直接使用LLM的有效输出,不再依赖预设的角色分配逻辑
|
|
|
341 |
|
342 |
if most_active_unknown:
|
343 |
final_map[most_active_unknown] = "Podcast Host"
|
344 |
+
|
345 |
+
print(f"LLM识别结果: {final_map}")
|
346 |
|
347 |
return final_map
|
348 |
|