OpenSTBench 0.2.0__py3-none-any.whl

openstbench/latency/agent.py

@@ -0,0 +1,108 @@
+from inspect import signature
+from typing import List
+
+from .basics import (
+    Action,
+    AgentStates,
+    EmptySegment,
+    Segment,
+    SpeechSegment,
+    TextSegment,
+)
+
+
+class GenericAgent:
+    """Base agent interface for latency evaluation."""
+
+    def __init__(self):
+        self.states = AgentStates()
+        self._pending_model_inference_time = None
+        self.reset()
+
+    def reset(self):
+        self.states.reset()
+        self._pending_model_inference_time = None
+
+    def record_model_inference_time(self, seconds: float) -> None:
+        if seconds is None:
+            return
+        seconds = float(seconds)
+        if seconds < 0:
+            return
+        if self._pending_model_inference_time is None:
+            self._pending_model_inference_time = 0.0
+        self._pending_model_inference_time += seconds
+
+    def consume_model_inference_time(self):
+        value = self._pending_model_inference_time
+        self._pending_model_inference_time = None
+        return value
+
+    def policy(self, states=None) -> Action:
+        raise NotImplementedError("Please implement policy().")
+
+    def push(self, segment: Segment, states=None) -> None:
+        (states or self.states).update_source(segment)
+
+    def pop(self, states=None) -> Segment:
+        state = states or self.states
+        if state.target_finished:
+            return EmptySegment(finished=True)
+
+        if len(signature(self.policy).parameters) > 0:
+            action = self.policy(state)
+        else:
+            action = self.policy()
+
+        if action.is_read():
+            return EmptySegment()
+
+        if isinstance(action.content, Segment):
+            segment = action.content
+            if action.finished is not None:
+                segment.finished = action.finished
+        elif isinstance(action.content, str):
+            segment = TextSegment(content=action.content, finished=action.finished)
+        elif isinstance(action.content, list):
+            segment = SpeechSegment(content=action.content, finished=action.finished)
+        else:
+            raise ValueError(f"Unknown content type: {type(action.content)}")
+
+        state.update_target(segment)
+        return segment
+
+    def pushpop(self, segment: Segment) -> Segment:
+        self.push(segment)
+        return self.pop()
+
+
+class AgentPipeline(GenericAgent):
+    """Sequentially compose multiple agents."""
+
+    def __init__(self, agents: List[GenericAgent]):
+        self.pipeline = agents
+        self.states = [agent.states for agent in agents]
+        self._pending_model_inference_time = None
+
+    def reset(self):
+        self._pending_model_inference_time = None
+        for agent in self.pipeline:
+            agent.reset()
+
+    def pushpop(self, segment: Segment) -> Segment:
+        current_input = segment
+        for index, agent in enumerate(self.pipeline):
+            if index == 0:
+                agent.push(current_input)
+                current_output = agent.pop()
+            else:
+                if not current_output.is_empty:
+                    agent.push(current_output)
+                current_output = agent.pop()
+
+            if hasattr(agent, "consume_model_inference_time"):
+                self.record_model_inference_time(agent.consume_model_inference_time())
+
+            if current_output.is_empty and index < len(self.pipeline) - 1:
+                return EmptySegment()
+        return current_output

openstbench/latency/basics.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass, field
+from typing import Union, List
+
+class Action:
+    def is_read(self) -> bool: raise NotImplementedError
+
+class ReadAction(Action):
+    def is_read(self) -> bool: return True
+    def __repr__(self): return "ReadAction"
+
+@dataclass
+class WriteAction(Action):
+    content: Union[str, List[float]]
+    finished: bool
+    def is_read(self) -> bool: return False
+
+@dataclass
+class Segment:
+    index: int = 0
+    content: list = field(default_factory=list)
+    finished: bool = False
+    is_empty: bool = False
+    data_type: str = None
+    config: dict = field(default_factory=dict)
+
+@dataclass
+class EmptySegment(Segment):
+    is_empty: bool = True
+
+@dataclass
+class TextSegment(Segment):
+    content: str = ""
+    data_type: str = "text"
+
+@dataclass
+class SpeechSegment(Segment):
+    sample_rate: int = 16000
+    data_type: str = "speech"
+
+class AgentStates:
+    def __init__(self): self.reset()
+    def reset(self):
+        self.source, self.target = [], []
+        self.source_finished = self.target_finished = False
+        self.upstream_states = []
+    
+    def update_source(self, segment: Segment):
+        self.source_finished = segment.finished
+        if not segment.is_empty:
+            if isinstance(segment, TextSegment): self.source.append(segment.content)
+            elif isinstance(segment, SpeechSegment): self.source += segment.content
+            
+    def update_target(self, segment: Segment):
+        self.target_finished = segment.finished
+        if not segment.is_empty and not self.target_finished:
+            if isinstance(segment, TextSegment): self.target.append(segment.content)
+            elif isinstance(segment, SpeechSegment): self.target += segment.content

openstbench/latency/cli.py

@@ -0,0 +1,257 @@
+import argparse
+import importlib.util
+import sys
+from pathlib import Path
+
+from tqdm import tqdm
+
+from .agent import GenericAgent
+from .instance import SpeechToTextInstance, SpeechToSpeechInstance
+from .metrics import SCORERS
+from .utils import Visualizer, materialize_s2s_alignment_artifacts, submit_slurm
+
+# --- 引入兄弟模块 (紧耦合) ---
+# 注意：使用相对导入或包绝对导入
+from ..translation_evaluator import TranslationEvaluator
+
+class LatencyEvaluator:
+    """同传延迟与质量评测器"""
+    def __init__(
+        self,
+        agent: GenericAgent,
+        segment_size=20,
+        latency_unit="word",
+        asr_fallback_for_s2s_alignment=True,
+        asr_model="medium",
+        asr_device=None,
+        alignment_acoustic_model="english_mfa",
+        alignment_dictionary_model="english_mfa",
+    ):
+        self.agent = agent
+        self.segment_size = segment_size
+        self.instances = {}
+        self.latency_unit = latency_unit
+        self.asr_fallback_for_s2s_alignment = asr_fallback_for_s2s_alignment
+        self.asr_model = asr_model
+        self.asr_device = asr_device
+        self.alignment_acoustic_model = alignment_acoustic_model
+        self.alignment_dictionary_model = alignment_dictionary_model
+
+    def run(self, source_files, ref_files, task="s2t", output_dir="./output", visualize=False):
+        Path(output_dir).mkdir(parents=True, exist_ok=True)
+        visualizer = Visualizer(output_dir) if visualize else None
+        
+        print(f"🚀 Running Latency Evaluation ({task}, {len(source_files)} samples)...")
+        for i, (src, ref) in tqdm(enumerate(zip(source_files, ref_files)), total=len(source_files)):
+            if task == "s2t":
+                ins = SpeechToTextInstance(i, src, ref, output_dir)
+            else:
+                ins = SpeechToSpeechInstance(i, src, ref, output_dir)
+            
+            self.agent.reset()
+            while not ins.finish_prediction:
+                s_in = ins.send_source(self.segment_size)
+                import time
+                t0 = time.perf_counter()
+                s_out = self.agent.pushpop(s_in)
+                t1 = time.perf_counter()
+                ins.add_inference_time(t1 - t0)
+                if hasattr(self.agent, "consume_model_inference_time"):
+                    ins.add_model_inference_time(self.agent.consume_model_inference_time())
+                ins.receive_prediction(s_out)
+            
+            self.instances[i] = ins
+            if visualizer: visualizer.plot(ins.summarize())
+        
+        # S2S 准备文本
+        if task == "s2s":
+            self._prepare_s2s_transcripts(output_dir)
+            
+        return self.instances
+
+    def _prepare_s2s_transcripts(self, output_dir):
+        materialize_s2s_alignment_artifacts(
+            self.instances,
+            output_dir,
+            unit=self.latency_unit,
+            asr_fallback=self.asr_fallback_for_s2s_alignment,
+            asr_model=self.asr_model,
+            asr_device=self.asr_device,
+        )
+
+        # 增加 show_all_metrics 开关，默认为 False
+    def compute_latency(self, computation_aware=False, output_dir="./output", show_all_metrics=False):
+        results = {}
+        # 1. 扫描并算出所有注册的 Scorer 成绩
+        for name, cls in SCORERS.items():
+            try:
+                if "Align" in name:
+                    scorer = cls(
+                        computation_aware=computation_aware,
+                        output_dir=output_dir,
+                        alignment_acoustic_model=self.alignment_acoustic_model,
+                        alignment_dictionary_model=self.alignment_dictionary_model,
+                    )
+                else:
+                    scorer = cls(computation_aware=computation_aware)
+                
+                score = scorer(self.instances)
+                key = f"{name}_CA" if computation_aware else name
+                results[key] = score
+            except Exception as e:
+                pass
+        
+        # 2. 智能化筛选最高级/最准确的一个版本展示给用户
+        cleaned_results = {}
+        
+        # [A] 评估: 第一声开口延迟 (StartOffset/ALAL)
+        so_key = "StartOffset_CA" if computation_aware else "StartOffset"
+        align_so_key = "StartOffset_SpeechAlign_CA" if computation_aware else "StartOffset_SpeechAlign"
+        
+        if results.get(align_so_key) is not None and results[align_so_key] > 0:
+            cleaned_results["First_Audio_Delay_(ALAL_ms)"] = results[align_so_key]
+        else:
+            cleaned_results["First_Audio_Delay_(ALAL_ms)"] = results.get(so_key, 0)
+            
+        # [B] 评估: 整句同传综合延迟 (标准版 ATD: 连同阅读/播放音频时间一起计算)
+        atd_key = "ATD_CA" if computation_aware else "ATD"
+        align_atd_key = "ATD_SpeechAlign_CA" if computation_aware else "ATD_SpeechAlign"
+        
+        if results.get(align_atd_key) is not None and results[align_atd_key] > 0:
+            cleaned_results["Overall_Translation_Delay_(ATD_ms)"] = results[align_atd_key]
+        else:
+            cleaned_results["Overall_Translation_Delay_(ATD_ms)"] = results.get(atd_key, 0)
+
+        # [C] 评估: 模型结单同传延迟 (纯净版 CustomATD: 剔除音频合成自身的物理时长，只看模型推断何时结束)
+        catd_key = "CustomATD_CA" if computation_aware else "CustomATD"
+        align_catd_key = "CustomATD_SpeechAlign_CA" if computation_aware else "CustomATD_SpeechAlign"
+        
+        if results.get(align_catd_key) is not None and results[align_catd_key] > 0:
+            cleaned_results["End_Action_Delay_(CustomATD_ms)"] = results[align_catd_key]
+        else:
+            cleaned_results["End_Action_Delay_(CustomATD_ms)"] = results.get(catd_key, 0)
+
+        # [D] 评估: 实时率指标 (RTF)
+        if "RTF" in results:
+            cleaned_results["Real_Time_Factor_(RTF)"] = results["RTF"]
+        elif "RTF_CA" in results:
+            cleaned_results["Real_Time_Factor_(RTF)"] = results["RTF_CA"]
+
+        if "ModelGenerateRTF" in results and results["ModelGenerateRTF"] is not None:
+            cleaned_results["Model_Generate_RTF"] = results["ModelGenerateRTF"]
+        elif "ModelGenerateRTF_CA" in results and results["ModelGenerateRTF_CA"] is not None:
+            cleaned_results["Model_Generate_RTF"] = results["ModelGenerateRTF_CA"]
+
+        # ================= 修改开始 =================
+        # 如果用户显式开启了展示开关，则将原始杂乱数据装入 "detailed_all_metrics"
+        if show_all_metrics:
+            cleaned_results["detailed_all_metrics"] = results
+        # ================= 修改结束 =================
+
+        return cleaned_results
+
+def load_agent_from_file(path, class_name):
+    """动态加载用户 Agent"""
+    spec = importlib.util.spec_from_file_location("user_agent", path)
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return getattr(mod, class_name)
+
+def main():
+    parser = argparse.ArgumentParser(description="OpenSTBench Latency Evaluator")
+    parser.add_argument("--source", required=True)
+    parser.add_argument("--target", default=None)
+    parser.add_argument("--output", default="./output")
+    parser.add_argument("--task", choices=["s2t", "s2s"], default="s2t")
+    parser.add_argument("--agent-script", required=True, help="Path to python file containing Agent class")
+    parser.add_argument("--agent-class", required=True, help="Name of the Agent class")
+    parser.add_argument("--segment-size", type=int, default=20)
+    parser.add_argument("--latency-unit", choices=["word", "char"], default="word")
+    parser.add_argument("--disable-asr-fallback", action="store_true")
+    parser.add_argument("--asr-model", default="medium")
+    parser.add_argument("--alignment-acoustic-model", default="english_mfa")
+    parser.add_argument("--alignment-dictionary-model", default="english_mfa")
+    parser.add_argument("--computation-aware", action="store_true")
+    parser.add_argument("--quality", action="store_true", help="Run Quality Evaluation (BLEU/WER) after latency")
+    parser.add_argument("--slurm", action="store_true")
+    args = parser.parse_args()
+
+    if args.slurm:
+        submit_slurm(args, __file__)
+
+    # 1. 加载数据
+    with open(args.source) as f: src = [l.strip() for l in f if l.strip()]
+    ref = [None]*len(src)
+    if args.target:
+        with open(args.target) as f: ref = [l.strip() for l in f if l.strip()]
+
+    # 2. 加载用户 Agent
+    AgentClass = load_agent_from_file(args.agent_script, args.agent_class)
+    agent = AgentClass()
+    
+    # 3. 运行 Latency 评测
+    evaluator = LatencyEvaluator(
+        agent,
+        args.segment_size,
+        latency_unit=args.latency_unit,
+        asr_fallback_for_s2s_alignment=not args.disable_asr_fallback,
+        asr_model=args.asr_model,
+        alignment_acoustic_model=args.alignment_acoustic_model,
+        alignment_dictionary_model=args.alignment_dictionary_model,
+    )
+    instances = evaluator.run(src, ref, args.task, args.output)
+    
+    # 4. 计算 Latency 指标
+    print("\n--- Latency Metrics ---")
+    scores = evaluator.compute_latency(False, args.output)
+    if args.computation_aware:
+        scores.update(evaluator.compute_latency(True, args.output))
+    
+    for k, v in scores.items():
+        print(f"{k:<25}: {v:.4f}")
+
+    # 5. [紧耦合] 运行 Quality 指标
+    if args.quality and args.target:
+        print("\n--- Quality Metrics (Integrated) ---")
+        
+        # 收集预测结果
+        # 注意: 对于 s2s, instances[i].prediction 应该是音频文件路径
+        predictions = [ins.get_prediction_content() for ins in instances.values()]
+        
+        if args.task == "s2t":
+            # [修改] 调用 TranslationEvaluator 算 BLEU
+            # 默认只开 BLEU 以保证速度
+            qual_eval = TranslationEvaluator(
+                use_bleu=True, 
+                use_chrf=False, 
+                use_comet=False, 
+                use_whisper=False
+            )
+            # 使用新的 evaluate_all 接口
+            q_res = qual_eval.evaluate_all(
+                target_text=predictions, 
+                reference=ref, 
+                source=src
+            )
+            print(q_res)
+            
+        elif args.task == "s2s":
+            # [修改] 调用 TranslationEvaluator 替代原 SpeechEvaluator
+            # 开启语音相关功能: use_wer=True, use_whisper=True
+            # predictions 是 wav 路径列表
+            speech_eval = TranslationEvaluator(
+                use_wer=True, 
+                use_whisper=True, 
+                whisper_model="tiny", # 用 tiny 快速出 WER
+                use_utmos=False       # 可选开启
+            )
+            # 使用 evaluate_all 接口传入 target_speech
+            s_res = speech_eval.evaluate_all(
+                target_speech=predictions, 
+                reference=ref,
+                source=src
+            )
+            print(s_res)
+
+if __name__ == "__main__":
+    main()

openstbench/latency/instance.py

@@ -0,0 +1,176 @@
+import time
+import math
+import soundfile
+from pathlib import Path
+from .basics import EmptySegment, SpeechSegment, Segment
+
+class Instance:
+    def __init__(self, index: int, source_path, reference=None, output_dir="./output"):
+        self.index = index
+        self.source_path = source_path
+        self.reference = reference
+        self.output_dir = Path(output_dir)
+        self.reset()
+
+    def reset(self):
+        self.step = 0
+        self.elapsed = [] 
+        self.delays = [] 
+        self.prediction_list = [] 
+        self.start_time = None
+        self.durations = [] # S2S 专用
+        self.finish_prediction = False
+        self.total_inference_time = 0.0 
+        self.total_model_inference_time = None
+        self.prediction_text = ""
+        self.prediction_text_source = "none"
+        self.target_units = []
+        self.target_unit_times_ms = []
+        self.alignment_mode = "none"
+        self.prediction_audio_path = None
+        self.source_chunk_end_times_ms = []
+
+    def summarize(self):
+        return {
+            "index": self.index,
+            "source": (self.source_path, ""),
+            "reference": self.reference,
+            "prediction": self.get_prediction_content(),
+            "prediction_length": len(self.prediction_list),
+            "delays": self.delays,
+            "elapsed": self.elapsed,
+            "durations": self.durations,
+            "prediction_text": self.prediction_text,
+            "prediction_text_source": self.prediction_text_source,
+            "alignment_mode": self.alignment_mode,
+        }
+    
+    def get_prediction_content(self): raise NotImplementedError
+
+    def add_model_inference_time(self, time_spent_in_seconds):
+        if time_spent_in_seconds is None:
+            return
+        time_spent_in_seconds = float(time_spent_in_seconds)
+        if time_spent_in_seconds < 0:
+            return
+        if self.total_model_inference_time is None:
+            self.total_model_inference_time = 0.0
+        self.total_model_inference_time += time_spent_in_seconds
+
+    def append_prediction_text(self, text: str, source: str = "native_transcript"):
+        text = (text or "").strip()
+        if not text:
+            return
+        if self.prediction_text:
+            sep = "" if self.prediction_text.endswith(" ") else " "
+            self.prediction_text = f"{self.prediction_text}{sep}{text}".strip()
+        else:
+            self.prediction_text = text
+        self.prediction_text_source = source
+    def get_prediction_raw(self): return self.prediction_list # 返回原始 list 供质量评测
+
+class SpeechToTextInstance(Instance):
+    def __init__(self, index, source_path, reference=None, output_dir="./output"):
+        data, sr = soundfile.read(source_path, dtype="float32")
+        self.samples = data.tolist()
+        self.sample_rate = sr
+        self.source_finished_reading = False
+        super().__init__(index, source_path, reference, output_dir)
+
+    def len_sample_to_ms(self, length): return length * 1000 / self.sample_rate
+
+    def send_source(self, segment_size=10):
+        if self.step == 0: self.start_time = time.time()
+        num_samples = math.ceil(segment_size / 1000 * self.sample_rate)
+        
+        if self.step < len(self.samples):
+            chunk = self.samples[self.step : self.step + num_samples]
+            is_finished = (self.step + num_samples >= len(self.samples))
+            self.step += len(chunk)
+            self.source_chunk_end_times_ms.append(self.len_sample_to_ms(self.step))
+            return SpeechSegment(content=chunk, sample_rate=self.sample_rate, finished=is_finished)
+        else:
+            self.source_finished_reading = True
+            return EmptySegment(finished=True)
+
+    def add_inference_time(self, time_spent_in_seconds: float):
+        """累加纯模型计算时间"""
+        self.total_inference_time += time_spent_in_seconds
+        
+    def receive_prediction(self, seg: Segment):
+        if self.finish_prediction or seg.is_empty:
+            self.finish_prediction = seg.finished
+            return
+        if not seg.content: return
+        if getattr(seg, "config", None):
+            self.add_model_inference_time(seg.config.get("model_inference_time"))
+        
+        current_time = time.time()
+        content_list = seg.content.strip().split()
+        self.prediction_list += content_list
+        self.append_prediction_text(seg.content, source="native_transcript")
+        
+        curr_delay = self.len_sample_to_ms(self.step)
+        curr_elapsed = curr_delay + (current_time - self.start_time) * 1000
+        
+        self.delays += [curr_delay] * len(content_list)
+        self.elapsed += [curr_elapsed] * len(content_list)
+        self.finish_prediction = seg.finished
+
+    def get_prediction_content(self): return " ".join(self.prediction_list)
+    
+    @property
+    def source_length(self): return self.len_sample_to_ms(len(self.samples))
+    @property
+    def reference_length(self): return len(self.reference.split()) if self.reference else 0
+
+class SpeechToSpeechInstance(SpeechToTextInstance):
+    def __init__(self, index, source_path, reference=None, output_dir="./output"):
+        super().__init__(index, source_path, reference, output_dir)
+        self.target_sample_rate = 16000
+
+    def receive_prediction(self, seg: Segment):
+        if self.finish_prediction or seg.is_empty:
+            self.finish_prediction = seg.finished
+            return
+        if not seg.content: return
+        if not self.start_time: self.start_time = time.time()
+        if getattr(seg, "config", None):
+            self.add_model_inference_time(seg.config.get("model_inference_time"))
+            transcript = seg.config.get("transcript")
+            if transcript:
+                source = seg.config.get("transcript_source", "native_transcript")
+                self.append_prediction_text(transcript, source=source)
+        
+        current_time = time.time()
+        duration_ms = len(seg.content) * 1000 / seg.sample_rate
+        self.target_sample_rate = seg.sample_rate
+        
+        self.prediction_list.append(seg.content)
+        self.durations.append(duration_ms)
+        
+        curr_delay = self.len_sample_to_ms(self.step)
+        curr_elapsed = curr_delay + (current_time - self.start_time) * 1000
+        
+        self.delays.append(curr_delay)
+        self.elapsed.append(curr_elapsed)
+        self.finish_prediction = seg.finished
+
+    def get_prediction_content(self):
+        # 返回 wav 路径
+        wav_dir = self.output_dir / "wavs"
+        wav_dir.mkdir(parents=True, exist_ok=True)
+        wav_path = wav_dir / f"{self.index}_pred.wav"
+        
+        flat_samples = [item for sublist in self.prediction_list for item in sublist]
+        if flat_samples:
+            soundfile.write(wav_path, flat_samples, self.target_sample_rate)
+        self.prediction_audio_path = str(wav_path.absolute())
+        return self.prediction_audio_path
+    
+    def get_prediction_raw(self):
+        # S2S 的原始预测结果是音频路径
+        return self.get_prediction_content()
+
+    @property 
+    def reference_length(self): return len(self.prediction_list)