OpenSTBench 0.2.0__py3-none-any.whl

openstbench/speaker_similarity_evaluator.py

@@ -0,0 +1,147 @@
+import torch
+import torchaudio
+import numpy as np
+import os
+import librosa
+from tqdm import tqdm
+
+# 如果你想用本地的 Resemblyzer，确保之前项目能 import 到
+try:
+    from resemblyzer import preprocess_wav, VoiceEncoder
+    RESEMBLYZER_AVAILABLE = True
+except ImportError:
+    RESEMBLYZER_AVAILABLE = False
+
+try:
+    from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector
+    TRANSFORMERS_AVAILABLE = True
+except ImportError:
+    TRANSFORMERS_AVAILABLE = False
+
+
+class SpeakerSimilarityEvaluator:
+    """
+    评估提取两段音频说话人特征（Speaker Embeddings）并计算余弦相似度的打分器。
+    """
+    def __init__(self, 
+                 model_type="wavlm", 
+                 device=None, 
+                 wavlm_model_path="microsoft/wavlm-base-plus-sv",
+                 resemblyzer_weights_path="pretrained.pt"):
+        """
+        初始化打分器，仅在初始化时加载一次模型。
+        
+        Args:
+            model_type (str): "wavlm" 或 "resemblyzer" 或 "both"。
+            device (str): "cuda", "mps" 或 "cpu"，如果为 None 则自动检测。
+            wavlm_model_path (str): WavLM 模型的路径或 HuggingFace ID。
+            resemblyzer_weights_path (str): Resemblyzer 的 pretrained.pt 权重路径。
+        """
+        self.device = device if device else (
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+        self.model_type = model_type.lower()
+        print(f"Initializing SpeakerSimilarityEvaluator with model(s): {self.model_type} on {self.device}")
+
+        self.wavlm_feature_extractor = None
+        self.wavlm_model = None
+        self.resemblyzer_encoder = None
+
+        if self.model_type in ["wavlm", "both"]:
+            if not TRANSFORMERS_AVAILABLE:
+                raise ImportError("Please install transformers to use WavLM: pip install transformers")
+            print(f"Loading WavLM from {wavlm_model_path}...")
+            self.wavlm_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wavlm_model_path)
+            self.wavlm_model = WavLMForXVector.from_pretrained(wavlm_model_path).to(self.device).eval()
+
+        if self.model_type in ["resemblyzer", "both"]:
+            if not RESEMBLYZER_AVAILABLE:
+                raise ImportError("Please ensure 'resemblyzer' module is available in your Python path.")
+            print("Loading Resemblyzer VoiceEncoder...")
+            # 兼容绝对路径或默认行为
+            if os.path.exists(resemblyzer_weights_path):
+                self.resemblyzer_encoder = VoiceEncoder(weights_fpath=resemblyzer_weights_path)
+            else:
+                self.resemblyzer_encoder = VoiceEncoder() # 内部会自动处理没传 path 的默认下载
+
+    @torch.no_grad()
+    def evaluate(self, ref_wav_path, synth_wav_path):
+        """
+        计算单对音频的相似度得分。
+        
+        Args:
+            ref_wav_path (str): 参考音频路径（Prompt / Target）
+            synth_wav_path (str): 你模型合成出来的音频路径
+            
+        Returns:
+            dict: 包含计算结果的字典，例如 {"wavlm": 0.85, "resemblyzer": 0.91}
+        """
+        results = {}
+
+        if self.model_type in ["wavlm", "both"]:
+            # WavLM 强制要求 16000 采样率
+            ref_waves_16k, _ = librosa.load(ref_wav_path, sr=16000)
+            synth_waves_16k, _ = librosa.load(synth_wav_path, sr=16000)
+
+            ref_inputs = self.wavlm_feature_extractor(
+                torch.tensor(ref_waves_16k), padding=True, return_tensors="pt", sampling_rate=16000
+            ).to(self.device)
+            synth_inputs = self.wavlm_feature_extractor(
+                torch.tensor(synth_waves_16k), padding=True, return_tensors="pt", sampling_rate=16000
+            ).to(self.device)
+
+            ref_embeddings = self.wavlm_model(**ref_inputs).embeddings
+            synth_embeddings = self.wavlm_model(**synth_inputs).embeddings
+
+            # L2 归一化后计算余弦相似度
+            ref_embeddings = torch.nn.functional.normalize(ref_embeddings, dim=-1)
+            synth_embeddings = torch.nn.functional.normalize(synth_embeddings, dim=-1)
+            similarity_wavlm = torch.nn.functional.cosine_similarity(ref_embeddings, synth_embeddings, dim=-1)
+            
+            results["wavlm_similarity"] = similarity_wavlm.item()
+
+        if self.model_type in ["resemblyzer", "both"]:
+            ref_wav_res = preprocess_wav(ref_wav_path)
+            synth_wav_res = preprocess_wav(synth_wav_path)
+
+            ref_embed_res = self.resemblyzer_encoder.embed_utterance(ref_wav_res)
+            synth_embed_res = self.resemblyzer_encoder.embed_utterance(synth_wav_res)
+            
+            similarity_res = np.inner(ref_embed_res, synth_embed_res)
+            results["resemblyzer_similarity"] = float(similarity_res)
+
+        return results
+
+    def evaluate_batch(self, ref_wav_paths, synth_wav_paths):
+        """
+        批量计算文件夹下的多对音频相似度。
+        
+        Args:
+            ref_wav_paths (list of str): 参考音频列表
+            synth_wav_paths (list of str): 合成音频列表（必须与参考列表一一对应）
+            
+        Returns:
+            dict: 包含各个模型平均分的字典
+        """
+        assert len(ref_wav_paths) == len(synth_wav_paths), "参考音频与合成音频数量必须对应！"
+        
+        batch_results = {"details": []}
+        wavlm_scores = []
+        res_scores = []
+
+        print("Evaluating speaker similarity...")
+        for ref_p, syn_p in tqdm(zip(ref_wav_paths, synth_wav_paths), total=len(ref_wav_paths)):
+            res = self.evaluate(ref_p, syn_p)
+            batch_results["details"].append({"ref": ref_p, "synth": syn_p, "score": res})
+            
+            if "wavlm_similarity" in res:
+                wavlm_scores.append(res["wavlm_similarity"])
+            if "resemblyzer_similarity" in res:
+                res_scores.append(res["resemblyzer_similarity"])
+
+        if wavlm_scores:
+            batch_results["average_wavlm_similarity"] = sum(wavlm_scores) / len(wavlm_scores)
+        if res_scores:
+            batch_results["average_resemblyzer_similarity"] = sum(res_scores) / len(res_scores)
+            
+        return batch_results

openstbench/speech_quality_evaluator.py

@@ -0,0 +1,171 @@
+import os
+import torch
+import torchaudio
+import jiwer
+from pathlib import Path
+from typing import List, Dict, Optional, Union
+from tqdm import tqdm
+
+# 复用现有的加载工具
+from .translation_evaluator import load_text_from_file_or_list, load_audio_from_folder
+
+try:
+    import whisper
+except ImportError:
+    whisper = None
+
+class SpeechQualityEvaluator:
+    """
+    语音质量与一致性评测器
+    专门用于：
+    1. UTMOS (音频自然度/听感质量)
+    2. WER/CER 一致性计算 (使用模型自己生成的文本与其生成的音频进行比对)
+    """
+    def __init__(self, 
+                 use_wer: bool = True,
+                 use_utmos: bool = True,
+                 whisper_model: str = "medium",
+                 utmos_model_path: Optional[str] = None,
+                 utmos_ckpt_path: Optional[str] = None,
+                 device: Optional[str] = None):
+        
+        self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
+        self.use_wer = use_wer
+        self.use_utmos = use_utmos
+        
+        self.whisper_model_name = whisper_model
+        self.utmos_path = utmos_model_path
+        self.utmos_ckpt = utmos_ckpt_path
+        
+        self.whisper_model = None
+        self.utmos_model = None
+        
+        if self.use_wer:
+            self.wer_transform = jiwer.Compose([
+                jiwer.ToLowerCase(),
+                jiwer.RemovePunctuation(),
+                jiwer.RemoveMultipleSpaces(),
+                jiwer.Strip(),
+                jiwer.RemoveEmptyStrings(),
+            ])
+
+    def _load_whisper(self):
+        if self.whisper_model is None and self.use_wer:
+            if not whisper:
+                print("⚠️ Whisper 未安装，跳过加载")
+                self.use_wer = False
+                return
+            print(f"⏳ 正在加载 Whisper ({self.whisper_model_name})...")
+            self.whisper_model = whisper.load_model(self.whisper_model_name, device=self.device)
+
+    def _load_utmos(self):
+        if self.utmos_model is None and self.use_utmos:
+            print("⏳ 正在加载 UTMOS 模型...")
+            try:
+                source = "github"
+                repo_or_dir = "tarepan/SpeechMOS"
+                if self.utmos_path and os.path.exists(self.utmos_path):
+                    source = "local"
+                    repo_or_dir = self.utmos_path
+
+                load_pretrained = (self.utmos_ckpt is None)
+                self.utmos_model = torch.hub.load(
+                    repo_or_dir, "utmos22_strong", source=source, trust_repo=True, pretrained=load_pretrained
+                )
+                
+                if self.utmos_ckpt and os.path.exists(self.utmos_ckpt):
+                    state_dict = torch.load(self.utmos_ckpt, map_location=self.device)
+                    self.utmos_model.load_state_dict(state_dict)
+                    
+                self.utmos_model.to(self.device).eval()
+            except Exception as e:
+                print(f"⚠️ UTMOS 模型加载失败: {e}")
+                self.use_utmos = False
+
+    def _transcribe(self, audio_paths: List[str]) -> List[str]:
+        self._load_whisper()
+        if not self.whisper_model:
+            return [""] * len(audio_paths)
+            
+        results = []
+        for path in tqdm(audio_paths, desc="🎙️ Whisper 转写中"):
+            res = self.whisper_model.transcribe(path, fp16=torch.cuda.is_available() and "cuda" in self.device)
+            results.append(res["text"].strip())
+        return results
+
+    def _compute_utmos(self, audio_paths: List[str]) -> float:
+        self._load_utmos()
+        if not self.utmos_model: return 0.0
+        
+        scores = []
+        target_sr = 16000
+        for path in tqdm(audio_paths, desc="🎧 计算 UTMOS"):
+            try:
+                wave, sr = torchaudio.load(path)
+                if sr != target_sr:
+                    wave = torchaudio.functional.resample(wave, sr, target_sr)
+                if wave.shape[0] > 1:
+                    wave = torch.mean(wave, dim=0, keepdim=True)
+                
+                wave = wave.to(self.device)
+                with torch.no_grad():
+                    s = self.utmos_model(wave, target_sr)
+                    scores.append(s.item())
+            except Exception as e:
+                print(f"⚠️ UTMOS Error on {path}: {e}")
+                
+        return sum(scores) / len(scores) if scores else 0.0
+
+    def _preprocess_for_wer(self, text: str, lang: str) -> str:
+        text = self.wer_transform(text)
+        if lang in ["zh", "ja", "ko"]:
+            text = text.replace(" ", "")
+            return " ".join(list(text))
+        return text
+
+    def evaluate_all(self, 
+                     target_audio: Union[List[str], str],
+                     target_text: Optional[Union[List[str], str]] = None,
+                     target_lang: str = "en") -> Dict[str, float]:
+        """
+        :param target_audio: 模型生成的音频路径列表或文件夹
+        :param target_text: 模型同步生成的文本，用作 WER 计算的参照
+        :param target_lang: 音频与文本的目标语言
+        """
+        results = {}
+        print(f"\n--- 开始语音质量评测 (Target Lang: {target_lang}) ---")
+        
+        # 加载目标音频
+        if isinstance(target_audio, str) and os.path.exists(target_audio) and os.path.isdir(target_audio):
+            audio_paths = load_audio_from_folder(target_audio)
+        else:
+            audio_paths = target_audio if isinstance(target_audio, list) else [target_audio]
+
+        # 1. 计算 UTMOS
+        if self.use_utmos:
+            print("   ➤ 计算 UTMOS (音频自然度)...")
+            results["UTMOS"] = round(self._compute_utmos(audio_paths), 4)
+
+        # 2. 计算 Consistency WER/CER
+        if self.use_wer:
+            if not target_text:
+                print("   ⚠️ 未提供 target_text (模型同时生成的文本)，无法计算文本-语音一致性，跳过 WER/CER。")
+            else:
+                texts = load_text_from_file_or_list(target_text, "target_text")
+                if len(texts) != len(audio_paths):
+                    raise ValueError(f"生成的文本行数 ({len(texts)}) 与 生成的音频数量 ({len(audio_paths)}) 不一致！")
+                
+                print("   ➤ 正在进行音频转写，准备计算一致性...")
+                asr_texts = self._transcribe(audio_paths)
+                
+                clean_refs = [self._preprocess_for_wer(t, target_lang) for t in texts]
+                clean_hyps = [self._preprocess_for_wer(t, target_lang) for t in asr_texts]
+                
+                if clean_refs:
+                    error_rate = jiwer.wer(clean_refs, clean_hyps)
+                    metric_name = "CER_Consistency" if target_lang in ['zh', 'ja', 'ko'] else "WER_Consistency"
+                    results[metric_name] = round(error_rate, 4)
+                else:
+                    print("   ⚠️ 提供的生成文本清洗后为空，无法计算分布错误率。")
+
+        return results

openstbench/translation_evaluator.py

@@ -0,0 +1,257 @@
+"""
+MultiMetric Eval - 文本翻译评测工具 (Text Metrics Only)
+"""
+import os
+import gc
+import json
+import numpy as np
+import sacrebleu
+import torch
+from typing import Dict, List, Optional, Union
+from pathlib import Path
+
+# ==================== 配置 ====================
+
+CACHE_PATHS = {
+    "huggingface": os.path.expanduser("~/.cache/huggingface/hub"),
+}
+
+# ==================== 可选依赖 ====================
+
+try:
+    from bleurt_pytorch import BleurtConfig, BleurtForSequenceClassification, BleurtTokenizer
+    HAS_BLEURT = True
+except ImportError:
+    HAS_BLEURT = False
+
+try:
+    from comet import download_model, load_from_checkpoint
+except ImportError:
+    download_model = None
+    load_from_checkpoint = None
+
+# ==================== 输入加载工具 ====================
+
+def load_text_from_file_or_list(input_data: Union[str, List[str]], name: str = "text") -> List[str]:
+    """
+    通用文本加载器
+    """
+    if isinstance(input_data, list):
+        return input_data
+    
+    if not isinstance(input_data, str):
+        raise ValueError(f"{name} 必须是 文件路径(str) 或 文本列表(List[str])")
+
+    path = Path(input_data)
+    if not path.exists():
+        raise FileNotFoundError(f"{name} 文件不存在: {input_data}")
+    
+    suffix = path.suffix.lower()
+    
+    if suffix == ".json":
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        
+        if isinstance(data, list):
+            if len(data) == 0: return []
+            if isinstance(data[0], str): return data
+            if isinstance(data[0], dict):
+                for key in ["target_text", "hypothesis", "text", "ref", "reference", "src", "source"]:
+                    if key in data[0]:
+                        return [item[key] for item in data]
+                raise ValueError(f"JSON 列表项中未找到常见文本字段")
+
+        if isinstance(data, dict):
+            for key in ["target_text", "hypothesis", "text", "ref", "reference", "src", "source"]:
+                if key in data:
+                    return data[key]
+            raise ValueError("JSON 字典中未找到常见文本字段")
+            
+        raise ValueError("不支持的 JSON 格式")
+    else:
+        with open(path, "r", encoding="utf-8") as f:
+            return [line.strip() for line in f if line.strip()]
+
+def load_audio_from_folder(folder_path: str, extensions: tuple = (".wav", ".mp3", ".flac")) -> List[str]:
+    """从文件夹加载音频文件路径（向后兼容暴露给 SpeechQualityEvaluator 使用）"""
+    folder = Path(folder_path)
+    if not folder.exists():
+        raise FileNotFoundError(f"文件夹不存在: {folder_path}")
+    audio_files = []
+    for ext in extensions:
+        audio_files.extend(folder.glob(f"*{ext}"))
+    audio_files = sorted(audio_files, key=lambda x: x.stem)
+    if not audio_files:
+        raise ValueError(f"文件夹中没有音频文件: {folder_path}")
+    return [str(f) for f in audio_files]
+
+
+# ==================== 评测器核心类 ====================
+
+DEFAULT_BLEURT_MODEL = "lucadiliello/BLEURT-20"
+
+class TranslationEvaluator:
+    """
+    文本侧翻译质量评测器：支持 文本翻译提取与对比 (BLEU, COMET, BLEURT...)
+    """
+
+    def __init__(self, 
+                 use_bleu: bool = True,
+                 use_chrf: bool = True,
+                 use_comet: bool = True,      
+                 use_bleurt: bool = False,    
+                 comet_model: str = "Unbabel/wmt22-comet-da",
+                 bleurt_path: Optional[str] = None,
+                 bleurt_model: Optional[str] = None,
+                 device: Optional[str] = None):
+        
+        self.use_bleu = use_bleu
+        self.use_chrf = use_chrf
+        self.use_comet = use_comet
+        self.use_bleurt = use_bleurt
+
+        if device is not None:
+            self.device = device
+        elif torch.cuda.is_available():
+            self.device = "cuda"
+        else:
+            self.device = "cpu"
+        
+        print(f"🚀 初始化 Translation Evaluator (纯文本侧) (设备: {self.device})")
+
+        # 加载语言计算模型
+        self.comet = None
+        if self.use_comet:
+            self.comet = self._load_comet(comet_model)
+
+        self.bleurt_model = None
+        self.bleurt_tokenizer = None
+        if self.use_bleurt:
+            self._load_bleurt(bleurt_path, bleurt_model)
+
+        print("✅ 翻译文本评测量度系统就绪！")
+
+    def __enter__(self): return self
+    def __exit__(self, exc_type, exc_val, exc_tb): 
+        self.cleanup()
+        return False
+
+    def cleanup(self):
+        for attr in ['comet', 'bleurt_model', 'bleurt_tokenizer']:
+            if hasattr(self, attr) and getattr(self, attr) is not None:
+                delattr(self, attr)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        gc.collect()
+
+    def _load_comet(self, model_name: str):
+        if not download_model:
+            print("⚠️ COMET 未安装，跳过")
+            return None
+        try:
+            cache = os.path.join(CACHE_PATHS["huggingface"], f"models--{model_name.replace('/', '--')}")
+            status = "[Local]" if os.path.exists(cache) else "[Online]"
+            print(f"⏳ {status} 加载 COMET: {model_name}")
+            model = load_from_checkpoint(download_model(model_name))
+            if self.device.startswith("cuda"):
+                model = model.to(self.device)
+            return model
+        except Exception as e:
+            print(f"❌ COMET 加载失败: {e}")
+            return None
+
+    def _load_bleurt(self, path: Optional[str], model_name: Optional[str]):
+        if not HAS_BLEURT:
+            print("⚠️ bleurt-pytorch 未安装，跳过加载")
+            return
+        
+        model_source = path if (path and os.path.exists(path)) else (model_name or DEFAULT_BLEURT_MODEL)
+        print(f"⏳ 加载 BLEURT: {model_source}")
+        
+        try:
+            self.bleurt_tokenizer = BleurtTokenizer.from_pretrained(model_source)
+            self.bleurt_model = BleurtForSequenceClassification.from_pretrained(model_source)
+            self.bleurt_model = self.bleurt_model.to(self.device).eval()
+        except Exception as e:
+            print(f"❌ BLEURT 加载失败: {e}")
+
+
+    def _get_bleu_tokenizer_name(self, lang: str) -> str:
+        if lang == 'zh': return 'zh'
+        elif lang == 'ja': return 'ja-mecab'
+        elif lang == 'ko': return 'ko-mecab'
+        else: return '13a'
+
+    def _compute_bleurt_score(self, references: List[str], candidates: List[str]) -> float:
+        all_scores = []
+        batch_size = 32
+        for i in range(0, len(references), batch_size):
+            br = references[i:i+batch_size]
+            bc = candidates[i:i+batch_size]
+            with torch.no_grad():
+                inputs = self.bleurt_tokenizer(br, bc, padding='longest', truncation=True, max_length=512, return_tensors='pt')
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                scores = self.bleurt_model(**inputs).logits.flatten().tolist()
+                all_scores.extend(scores)
+        return float(np.mean(all_scores))
+
+    def evaluate_all(
+        self,
+        reference: Union[List[str], str],
+        target_text: Union[List[str], str],
+        source: Optional[Union[List[str], str]] = None,
+        target_lang: str = "en"
+    ) -> Dict[str, float]:
+        """
+        Args:
+            reference: 数据集自带的真实标准基准翻译参考 (List 或 文件路径)
+            target_text: 翻译模型输出的直接文本翻译结果 (List 或 文件路径)
+            source: 选填, 语言原始文本 (COMET强制依赖)
+            target_lang: 影响 BLEU的 Tokenizer 选择
+        """
+        results = {}
+        print(f"\n--- 开始文本翻译质量评测 (Target Lang: {target_lang}) ---")
+
+        final_ref = load_text_from_file_or_list(reference, "Reference")
+        final_text = load_text_from_file_or_list(target_text, "Target Text")
+        
+        final_src = None
+        if source:
+            final_src = load_text_from_file_or_list(source, "Source")
+            if len(final_src) != len(final_ref):
+                raise ValueError("Source 与 Reference 数量不一致")
+        
+        if len(final_text) != len(final_ref):
+            raise ValueError("Target Text 与 Reference 数量不一致")
+
+        # 1. sacreBLEU
+        if self.use_bleu:
+            tokenizer_name = self._get_bleu_tokenizer_name(target_lang)
+            try:
+                results["sacreBLEU"] = sacrebleu.corpus_bleu(final_text, [final_ref], tokenize=tokenizer_name).score
+            except Exception as e:
+                results["sacreBLEU"] = -1.0
+
+        # 2. chrF++
+        if self.use_chrf:
+            try:
+                results["chrF++"] = sacrebleu.corpus_chrf(final_text, [final_ref], word_order=2).score
+            except: results["chrF++"] = -1.0
+
+        # 3. BLEURT
+        if self.use_bleurt and self.bleurt_model:
+            try:
+                results["BLEURT"] = self._compute_bleurt_score(final_ref, final_text)
+            except Exception as e:
+                results["BLEURT"] = -1.0
+
+        # 4. COMET
+        if self.use_comet and self.comet and final_src:
+            try:
+                data = [{"src": s, "mt": t, "ref": r} for s, t, r in zip(final_src, final_text, final_ref)]
+                gpus = 1 if self.device.startswith("cuda") else 0
+                results["COMET"] = self.comet.predict(data, batch_size=8, gpus=gpus).system_score
+            except Exception as e:
+                results["COMET"] = -1.0
+        
+        return {k: round(v, 4) if v >= 0 else v for k, v in results.items()}