feat(13-07): integrate SemanticSegmentationEngine into SRTProcessor
- 添加 enable_semantic_segmentation 参数到 SRTProcessor.__init__ - 启用时使用 SemanticSegmentationEngine，禁用时使用 SemanticSegmenter - AI 不可用时自动降级到规则分段 - 更新 _generate_markdown() 支持带标题的 VideoSegment 输出 - 调整步骤顺序：先加载元数据，再进行语义分段 - 更新 UnifiedSegmentationStrategy.segment() 返回 List[VideoSegment] - 添加代理方法到 SemanticSegmentationEngine 以保持向后兼容 - 创建全面的集成测试（9 个测试用例） - 更新现有语义引擎测试以适配新 API Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -17,7 +17,7 @@ import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, Any, List, Optional
|
||||
from typing import Dict, Any, List, Optional, Union
|
||||
|
||||
from kb_workflow.pipeline.config import get_settings
|
||||
from kb_workflow.pipeline.processors.base import BaseProcessor
|
||||
@@ -48,19 +48,45 @@ class SRTProcessor(BaseProcessor):
|
||||
def __init__(
|
||||
self,
|
||||
state_manager: Optional[WorkflowStateManager] = None,
|
||||
enable_correction: bool = True
|
||||
enable_correction: bool = True,
|
||||
enable_semantic_segmentation: bool = False
|
||||
):
|
||||
"""初始化 SRT 处理器
|
||||
|
||||
Args:
|
||||
state_manager: 状态管理器
|
||||
enable_correction: 是否启用 ASR 纠错（默认 True）
|
||||
enable_semantic_segmentation: 是否启用 LLM 语义分段（默认 False）
|
||||
"""
|
||||
super().__init__(state_manager)
|
||||
self.enable_correction = enable_correction
|
||||
self.segmenter = SemanticSegmenter()
|
||||
self.enable_semantic_segmentation = enable_semantic_segmentation
|
||||
self.segmenter = None
|
||||
self.corrector = None
|
||||
|
||||
# 初始化分段器
|
||||
if enable_semantic_segmentation:
|
||||
try:
|
||||
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
|
||||
from kb_workflow.srt.semantic_engine import UnifiedSegmentationStrategy
|
||||
from kb_workflow.deep_analyzer import SimpleAIClient
|
||||
|
||||
# 创建 AI 客户端
|
||||
ai_client = SimpleAIClient()
|
||||
if ai_client.api_key:
|
||||
# 创建统一策略，传入 AI 客户端
|
||||
strategy = UnifiedSegmentationStrategy(ai_client=ai_client)
|
||||
self.segmenter = SemanticSegmentationEngine(strategy=strategy)
|
||||
logger.info("Semantic segmentation enabled")
|
||||
else:
|
||||
logger.info("AI API key not configured, using rule-based segmenter")
|
||||
self.segmenter = SemanticSegmenter()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize semantic engine: {e}, using rule-based segmenter")
|
||||
self.segmenter = SemanticSegmenter()
|
||||
else:
|
||||
self.segmenter = SemanticSegmenter()
|
||||
|
||||
# 延迟初始化纠错器
|
||||
if enable_correction:
|
||||
try:
|
||||
@@ -102,27 +128,53 @@ class SRTProcessor(BaseProcessor):
|
||||
segments = parse_srt_file(Path(input_data.srt_file))
|
||||
logger.info(f"Parsed {len(segments)} segments")
|
||||
|
||||
# Step 2: 语义分段
|
||||
# Step 2: 加载元数据（需要在分段之前加载，因为 LLM 分段需要元数据）
|
||||
metadata = self._load_metadata(input_data.metadata_file)
|
||||
|
||||
# Step 3: 语义分段
|
||||
logger.info("Segmenting subtitles into paragraphs")
|
||||
paragraphs = self.segmenter.segment(segments)
|
||||
if self.enable_semantic_segmentation:
|
||||
# 使用 LLM 语义分段
|
||||
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
|
||||
if isinstance(self.segmenter, SemanticSegmentationEngine):
|
||||
# 返回 VideoSegment 列表（包含标题）
|
||||
paragraphs = self.segmenter.segment(segments, metadata=metadata)
|
||||
else:
|
||||
# 规则分段返回纯文本列表
|
||||
paragraphs = self.segmenter.segment(segments)
|
||||
else:
|
||||
# 规则分段返回纯文本列表
|
||||
paragraphs = self.segmenter.segment(segments)
|
||||
logger.info(f"Generated {len(paragraphs)} paragraphs")
|
||||
|
||||
# Step 3: ASR 纠错（可选）
|
||||
# Step 4: ASR 纠错（可选）
|
||||
corrected = False
|
||||
if self.corrector and paragraphs:
|
||||
logger.info("Applying ASR correction")
|
||||
corrected_paragraphs = []
|
||||
for para in paragraphs:
|
||||
corrected_para = self.corrector.correct(para)
|
||||
# 处理 VideoSegment 或纯文本
|
||||
if hasattr(para, 'content'):
|
||||
# VideoSegment 对象
|
||||
content = para.content
|
||||
corrected_content = self.corrector.correct(content)
|
||||
# 创建新的 VideoSegment 对象（假设不可变）
|
||||
from kb_workflow.srt.semantic_engine import VideoSegment
|
||||
corrected_para = VideoSegment(
|
||||
title=para.title,
|
||||
content=corrected_content,
|
||||
start_index=para.start_index,
|
||||
end_index=para.end_index
|
||||
)
|
||||
else:
|
||||
# 纯文本
|
||||
corrected_para = self.corrector.correct(para)
|
||||
corrected_paragraphs.append(corrected_para)
|
||||
if corrected_para != para:
|
||||
corrected = True
|
||||
paragraphs = corrected_paragraphs
|
||||
logger.info(f"ASR correction applied: {corrected}")
|
||||
|
||||
# Step 4: 加载元数据
|
||||
metadata = self._load_metadata(input_data.metadata_file)
|
||||
|
||||
# Step 5: 生成 Markdown
|
||||
logger.info(f"Generating Markdown: {input_data.output_file}")
|
||||
tags = self._extract_tags(paragraphs)
|
||||
@@ -133,6 +185,12 @@ class SRTProcessor(BaseProcessor):
|
||||
tags=tags
|
||||
)
|
||||
|
||||
# 转换 VideoSegment 为纯文本（用于 SRTOutput）
|
||||
if paragraphs and hasattr(paragraphs[0], 'content'):
|
||||
paragraphs_text = [para.content if hasattr(para, 'content') else para for para in paragraphs]
|
||||
else:
|
||||
paragraphs_text = paragraphs
|
||||
|
||||
# 记录完成状态
|
||||
if self.state_manager:
|
||||
state = self.state_manager.load()
|
||||
@@ -144,7 +202,7 @@ class SRTProcessor(BaseProcessor):
|
||||
return SRTOutput(
|
||||
success=True,
|
||||
output_file=str(input_data.output_file),
|
||||
paragraphs=paragraphs,
|
||||
paragraphs=paragraphs_text,
|
||||
tags=tags,
|
||||
corrected=corrected
|
||||
)
|
||||
@@ -190,14 +248,14 @@ class SRTProcessor(BaseProcessor):
|
||||
|
||||
return raw_metadata
|
||||
|
||||
def _extract_tags(self, paragraphs: List[str]) -> List[str]:
|
||||
def _extract_tags(self, paragraphs: Union[List[str], List[Any]]) -> List[str]:
|
||||
"""从段落中提取标签
|
||||
|
||||
简单实现：不使用 LLM，只返回默认标签。
|
||||
未来可以增强为 LLM 提取。
|
||||
|
||||
Args:
|
||||
paragraphs: 段落列表
|
||||
paragraphs: 段落列表（可以是纯文本或 VideoSegment 对象）
|
||||
|
||||
Returns:
|
||||
标签列表
|
||||
@@ -210,7 +268,7 @@ class SRTProcessor(BaseProcessor):
|
||||
self,
|
||||
output_file: Path,
|
||||
metadata: Dict[str, Any],
|
||||
paragraphs: List[str],
|
||||
paragraphs: Union[List[str], List[Any]], # List[str] 或 List[VideoSegment]
|
||||
tags: List[str]
|
||||
) -> None:
|
||||
"""生成 Markdown 文件
|
||||
@@ -230,12 +288,29 @@ class SRTProcessor(BaseProcessor):
|
||||
# 构建内容
|
||||
content_parts = [frontmatter, ""]
|
||||
|
||||
# 直接添加内容，不生成摘要部分
|
||||
content_parts.append("## 内容")
|
||||
content_parts.append("")
|
||||
for para in paragraphs:
|
||||
content_parts.append(para)
|
||||
# 检查是否是 VideoSegment 类型（LLM 分段结果）
|
||||
from kb_workflow.srt.semantic_engine import VideoSegment
|
||||
is_semantic_segmentation = (
|
||||
paragraphs and
|
||||
hasattr(paragraphs[0], 'title') and
|
||||
hasattr(paragraphs[0], 'content')
|
||||
)
|
||||
|
||||
if is_semantic_segmentation:
|
||||
# LLM 语义分段：包含标题和内容
|
||||
for seg in paragraphs:
|
||||
if seg.title:
|
||||
content_parts.append(f"## {seg.title}")
|
||||
content_parts.append("")
|
||||
content_parts.append(seg.content)
|
||||
content_parts.append("")
|
||||
else:
|
||||
# 规则分段：纯文本
|
||||
content_parts.append("## 内容")
|
||||
content_parts.append("")
|
||||
for para in paragraphs:
|
||||
content_parts.append(para)
|
||||
content_parts.append("")
|
||||
|
||||
# 写入文件
|
||||
output_file.write_text('\n'.join(content_parts), encoding='utf-8')
|
||||
|
||||
Reference in New Issue
Block a user