OpenSTBench 1.1.0__tar.gz → 1.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (27) hide show
  1. {openstbench-1.1.0/src/OpenSTBench.egg-info → openstbench-1.2.2}/PKG-INFO +4 -1
  2. {openstbench-1.1.0 → openstbench-1.2.2}/pyproject.toml +4 -3
  3. {openstbench-1.1.0 → openstbench-1.2.2/src/OpenSTBench.egg-info}/PKG-INFO +4 -1
  4. {openstbench-1.1.0 → openstbench-1.2.2}/src/OpenSTBench.egg-info/SOURCES.txt +1 -0
  5. {openstbench-1.1.0 → openstbench-1.2.2}/src/OpenSTBench.egg-info/requires.txt +3 -0
  6. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/__init__.py +4 -2
  7. openstbench-1.2.2/src/openstbench/language_policy.py +91 -0
  8. openstbench-1.2.2/src/openstbench/latency/__init__.py +39 -0
  9. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/speech_quality_evaluator.py +185 -181
  10. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/translation_evaluator.py +14 -7
  11. openstbench-1.1.0/src/openstbench/latency/__init__.py +0 -14
  12. {openstbench-1.1.0 → openstbench-1.2.2}/LICENSE +0 -0
  13. {openstbench-1.1.0 → openstbench-1.2.2}/README.md +0 -0
  14. {openstbench-1.1.0 → openstbench-1.2.2}/setup.cfg +0 -0
  15. {openstbench-1.1.0 → openstbench-1.2.2}/src/OpenSTBench.egg-info/dependency_links.txt +0 -0
  16. {openstbench-1.1.0 → openstbench-1.2.2}/src/OpenSTBench.egg-info/top_level.txt +0 -0
  17. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/_model_loading.py +0 -0
  18. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/emotion_evaluator.py +0 -0
  19. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/latency/agent.py +0 -0
  20. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/latency/basics.py +0 -0
  21. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/latency/cli.py +0 -0
  22. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/latency/instance.py +0 -0
  23. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/latency/metrics.py +0 -0
  24. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/latency/utils.py +0 -0
  25. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/paralinguistic_evaluator.py +0 -0
  26. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/speaker_similarity_evaluator.py +0 -0
  27. {openstbench-1.1.0 → openstbench-1.2.2}/src/openstbench/temporal_consistency_evaluator.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: OpenSTBench
3
- Version: 1.1.0
3
+ Version: 1.2.2
4
4
  Summary: Unified multidimensional evaluation toolkit for S2TT and S2ST systems in offline and streaming speech translation settings
5
5
  Author-email: Yanjie An <691476922@qq.com>
6
6
  License: MIT
@@ -17,6 +17,7 @@ Classifier: Programming Language :: Python :: 3
17
17
  Requires-Python: <3.11,>=3.9
18
18
  Description-Content-Type: text/markdown
19
19
  License-File: LICENSE
20
+ Requires-Dist: setuptools==80.9.0
20
21
  Requires-Dist: torch<2.7,>=2.6
21
22
  Requires-Dist: torchaudio<2.7,>=2.6
22
23
  Requires-Dist: numpy
@@ -36,6 +37,7 @@ Provides-Extra: whisper
36
37
  Requires-Dist: openai-whisper; extra == "whisper"
37
38
  Provides-Extra: speech-quality
38
39
  Requires-Dist: jiwer; extra == "speech-quality"
40
+ Requires-Dist: opencc-python-reimplemented; extra == "speech-quality"
39
41
  Provides-Extra: emotion
40
42
  Requires-Dist: funasr; extra == "emotion"
41
43
  Requires-Dist: modelscope; extra == "emotion"
@@ -45,6 +47,7 @@ Provides-Extra: all
45
47
  Requires-Dist: unbabel-comet>=2.2.2; extra == "all"
46
48
  Requires-Dist: openai-whisper; extra == "all"
47
49
  Requires-Dist: jiwer; extra == "all"
50
+ Requires-Dist: opencc-python-reimplemented; extra == "all"
48
51
  Requires-Dist: funasr; extra == "all"
49
52
  Requires-Dist: modelscope; extra == "all"
50
53
  Requires-Dist: transformers; extra == "all"
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "OpenSTBench"
7
- version = "1.1.0"
7
+ version = "1.2.2"
8
8
  description = "Unified multidimensional evaluation toolkit for S2TT and S2ST systems in offline and streaming speech translation settings"
9
9
  readme = {file = "README.md", content-type = "text/markdown"}
10
10
  requires-python = ">=3.9,<3.11"
@@ -29,6 +29,7 @@ classifiers = [
29
29
  ]
30
30
 
31
31
  dependencies = [
32
+ "setuptools==80.9.0",
32
33
  "torch>=2.6,<2.7",
33
34
  "torchaudio>=2.6,<2.7",
34
35
  "numpy",
@@ -53,10 +54,10 @@ Documentation = "https://github.com/sjtuayj/OpenSTBench#readme"
53
54
  [project.optional-dependencies]
54
55
  comet = ["unbabel-comet>=2.2.2"]
55
56
  whisper = ["openai-whisper"]
56
- speech_quality = ["jiwer"]
57
+ speech_quality = ["jiwer", "opencc-python-reimplemented"]
57
58
  emotion = ["funasr", "modelscope"]
58
59
  paralinguistics = ["transformers"]
59
- all = ["unbabel-comet>=2.2.2", "openai-whisper", "jiwer", "funasr", "modelscope", "transformers"]
60
+ all = ["unbabel-comet>=2.2.2", "openai-whisper", "jiwer", "opencc-python-reimplemented", "funasr", "modelscope", "transformers"]
60
61
 
61
62
  [tool.setuptools]
62
63
  package-dir = {"" = "src"}
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: OpenSTBench
3
- Version: 1.1.0
3
+ Version: 1.2.2
4
4
  Summary: Unified multidimensional evaluation toolkit for S2TT and S2ST systems in offline and streaming speech translation settings
5
5
  Author-email: Yanjie An <691476922@qq.com>
6
6
  License: MIT
@@ -17,6 +17,7 @@ Classifier: Programming Language :: Python :: 3
17
17
  Requires-Python: <3.11,>=3.9
18
18
  Description-Content-Type: text/markdown
19
19
  License-File: LICENSE
20
+ Requires-Dist: setuptools==80.9.0
20
21
  Requires-Dist: torch<2.7,>=2.6
21
22
  Requires-Dist: torchaudio<2.7,>=2.6
22
23
  Requires-Dist: numpy
@@ -36,6 +37,7 @@ Provides-Extra: whisper
36
37
  Requires-Dist: openai-whisper; extra == "whisper"
37
38
  Provides-Extra: speech-quality
38
39
  Requires-Dist: jiwer; extra == "speech-quality"
40
+ Requires-Dist: opencc-python-reimplemented; extra == "speech-quality"
39
41
  Provides-Extra: emotion
40
42
  Requires-Dist: funasr; extra == "emotion"
41
43
  Requires-Dist: modelscope; extra == "emotion"
@@ -45,6 +47,7 @@ Provides-Extra: all
45
47
  Requires-Dist: unbabel-comet>=2.2.2; extra == "all"
46
48
  Requires-Dist: openai-whisper; extra == "all"
47
49
  Requires-Dist: jiwer; extra == "all"
50
+ Requires-Dist: opencc-python-reimplemented; extra == "all"
48
51
  Requires-Dist: funasr; extra == "all"
49
52
  Requires-Dist: modelscope; extra == "all"
50
53
  Requires-Dist: transformers; extra == "all"
@@ -9,6 +9,7 @@ src/OpenSTBench.egg-info/top_level.txt
9
9
  src/openstbench/__init__.py
10
10
  src/openstbench/_model_loading.py
11
11
  src/openstbench/emotion_evaluator.py
12
+ src/openstbench/language_policy.py
12
13
  src/openstbench/paralinguistic_evaluator.py
13
14
  src/openstbench/speaker_similarity_evaluator.py
14
15
  src/openstbench/speech_quality_evaluator.py
@@ -1,3 +1,4 @@
1
+ setuptools==80.9.0
1
2
  torch<2.7,>=2.6
2
3
  torchaudio<2.7,>=2.6
3
4
  numpy
@@ -16,6 +17,7 @@ resemblyzer
16
17
  unbabel-comet>=2.2.2
17
18
  openai-whisper
18
19
  jiwer
20
+ opencc-python-reimplemented
19
21
  funasr
20
22
  modelscope
21
23
  transformers
@@ -32,6 +34,7 @@ transformers
32
34
 
33
35
  [speech_quality]
34
36
  jiwer
37
+ opencc-python-reimplemented
35
38
 
36
39
  [whisper]
37
40
  openai-whisper
@@ -9,7 +9,7 @@ actually accessed.
9
9
  from importlib import import_module
10
10
  from typing import Dict, Tuple
11
11
 
12
- __version__ = "1.1.0"
12
+ __version__ = "1.2.2"
13
13
 
14
14
  __all__ = [
15
15
  "TranslationEvaluator",
@@ -35,6 +35,8 @@ __all__ = [
35
35
  "AgentPipeline",
36
36
  "ReadAction",
37
37
  "WriteAction",
38
+ "TextSegment",
39
+ "SpeechSegment",
38
40
  "load_text_from_file_or_list",
39
41
  "load_audio_from_folder",
40
42
  # "Dataset",
@@ -87,7 +89,7 @@ _EXPORT_SPECS: Tuple[Tuple[str, Tuple[str, ...]], ...] = (
87
89
  # ),
88
90
  # ),
89
91
  ("latency.agent", ("AgentPipeline", "GenericAgent")),
90
- ("latency.basics", ("ReadAction", "WriteAction")),
92
+ ("latency.basics", ("ReadAction", "WriteAction", "TextSegment", "SpeechSegment")),
91
93
  ("latency.cli", ("LatencyEvaluator",)),
92
94
  )
93
95
 
@@ -0,0 +1,91 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Optional
4
+
5
+
6
+ LANGUAGE_POLICY_VERSION = "2026-07-fleurs-wer-cer-v1"
7
+
8
+ _LANGUAGE_ALIASES = {
9
+ "cmn": "zh",
10
+ "zho": "zh",
11
+ "chi": "zh",
12
+ "zh": "zh",
13
+ "yue": "zh",
14
+ "eng": "en",
15
+ "en": "en",
16
+ "fra": "fr",
17
+ "fre": "fr",
18
+ "fr": "fr",
19
+ "deu": "de",
20
+ "ger": "de",
21
+ "de": "de",
22
+ "spa": "es",
23
+ "es": "es",
24
+ "jpn": "ja",
25
+ "jp": "ja",
26
+ "ja": "ja",
27
+ "kor": "ko",
28
+ "ko": "ko",
29
+ "tha": "th",
30
+ "th": "th",
31
+ "lao": "lo",
32
+ "lo": "lo",
33
+ "khm": "km",
34
+ "km": "km",
35
+ "mya": "my",
36
+ "bur": "my",
37
+ "my": "my",
38
+ "bod": "bo",
39
+ "tib": "bo",
40
+ "bo": "bo",
41
+ "dzo": "dz",
42
+ "dz": "dz",
43
+ }
44
+
45
+ _CER_LANGUAGES = {
46
+ "zh",
47
+ "ja",
48
+ "ko",
49
+ "th",
50
+ "lo",
51
+ "km",
52
+ "my",
53
+ "bo",
54
+ "dz",
55
+ }
56
+
57
+
58
+ def normalize_language_code(language: Optional[str]) -> str:
59
+ text = str(language or "").strip().lower()
60
+ if not text:
61
+ return ""
62
+ if text.startswith("<|") and text.endswith("|>"):
63
+ text = text[2:-2]
64
+ text = text.replace("-", "_")
65
+ base = text.split("_", 1)[0]
66
+ return _LANGUAGE_ALIASES.get(base, base)
67
+
68
+
69
+ def speech_consistency_unit(language: Optional[str]) -> str:
70
+ normalized = normalize_language_code(language)
71
+ return "cer" if normalized in _CER_LANGUAGES else "wer"
72
+
73
+
74
+ def whisper_language_code(language: Optional[str]) -> Optional[str]:
75
+ normalized = normalize_language_code(language)
76
+ return normalized or None
77
+
78
+
79
+ def tokenization_metadata(language: Optional[str]) -> dict:
80
+ normalized = normalize_language_code(language)
81
+ warnings = []
82
+ if not normalized:
83
+ warnings.append("empty_language_code_defaulted_to_wer")
84
+ return {
85
+ "language_policy_version": LANGUAGE_POLICY_VERSION,
86
+ "target_language": str(language or ""),
87
+ "normalized_language": normalized,
88
+ "speech_consistency_unit": speech_consistency_unit(language),
89
+ "whisper_language": whisper_language_code(language),
90
+ "warnings": warnings,
91
+ }
@@ -0,0 +1,39 @@
1
+ # This file contains code adapted from SimulEval:
2
+ # https://github.com/facebookresearch/SimulEval
3
+ #
4
+ # SimulEval is licensed under the Creative Commons Attribution-ShareAlike
5
+ # 4.0 International License (CC BY-SA 4.0):
6
+ # https://creativecommons.org/licenses/by-sa/4.0/
7
+ #
8
+ # The adapted portions in this file are distributed under CC BY-SA 4.0.
9
+ # Modifications were made for OpenSTBench.
10
+
11
+ from .agent import GenericAgent, AgentPipeline
12
+ from .basics import ReadAction, WriteAction, TextSegment, SpeechSegment
13
+
14
+ __all__ = [
15
+ "GenericAgent",
16
+ "AgentPipeline",
17
+ "ReadAction",
18
+ "WriteAction",
19
+ "TextSegment",
20
+ "SpeechSegment",
21
+ "LatencyEvaluator",
22
+ "register",
23
+ ]
24
+
25
+
26
+ def __getattr__(name):
27
+ if name == "LatencyEvaluator":
28
+ from .cli import LatencyEvaluator
29
+
30
+ return LatencyEvaluator
31
+ if name == "register":
32
+ from .metrics import register
33
+
34
+ return register
35
+ raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
36
+
37
+
38
+ def __dir__():
39
+ return sorted(set(globals()) | set(__all__))
@@ -1,202 +1,206 @@
1
- import os
2
- import torch
3
- import torchaudio
4
- import jiwer
5
- from pathlib import Path
6
- from typing import List, Dict, Optional, Union
7
- from tqdm import tqdm
8
-
9
- # Reuse existing loading utilities
1
+ import os
2
+ import torch
3
+ import torchaudio
4
+ import jiwer
5
+ from pathlib import Path
6
+ from typing import List, Dict, Optional, Union
7
+ from tqdm import tqdm
8
+
9
+ # Reuse existing loading utilities
10
10
  from ._model_loading import resolve_pretrained_source
11
+ from .language_policy import normalize_language_code, speech_consistency_unit, whisper_language_code
11
12
  from .translation_evaluator import load_text_from_file_or_list, load_audio_from_folder
12
-
13
+
13
14
  try:
14
15
  import whisper
15
16
  except ImportError:
16
17
  whisper = None
17
18
 
18
- class SpeechQualityEvaluator:
19
- DEFAULT_WHISPER_MODEL = "medium"
20
- """
21
- Speech Quality and Consistency Evaluator.
22
- Specifically designed for:
23
- 1. UTMOS (Audio naturalness / perceived quality).
24
- 2. WER/CER Consistency Calculation (compares the model's generated text with its generated audio).
25
- """
26
- def __init__(self,
27
- use_wer: bool = True,
28
- use_utmos: bool = True,
29
- whisper_model: str = DEFAULT_WHISPER_MODEL,
30
- whisper_language: Optional[str] = None,
31
- utmos_model_path: Optional[str] = None,
32
- utmos_ckpt_path: Optional[str] = None,
33
- device: Optional[str] = None):
34
-
35
- self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
36
- self.use_wer = use_wer
37
- self.use_utmos = use_utmos
38
-
19
+ try:
20
+ from opencc import OpenCC
21
+ except ImportError:
22
+ OpenCC = None
23
+
24
+ class SpeechQualityEvaluator:
25
+ DEFAULT_WHISPER_MODEL = "medium"
26
+ """
27
+ Speech Quality and Consistency Evaluator.
28
+ Specifically designed for:
29
+ 1. UTMOS (Audio naturalness / perceived quality).
30
+ 2. WER/CER Consistency Calculation (compares the model's generated text with its generated audio).
31
+ """
32
+ def __init__(self,
33
+ use_wer: bool = True,
34
+ use_utmos: bool = True,
35
+ whisper_model: str = DEFAULT_WHISPER_MODEL,
36
+ whisper_language: Optional[str] = None,
37
+ utmos_model_path: Optional[str] = None,
38
+ utmos_ckpt_path: Optional[str] = None,
39
+ device: Optional[str] = None):
40
+
41
+ self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
42
+ self.use_wer = use_wer
43
+ self.use_utmos = use_utmos
44
+
39
45
  self.whisper_model_name = whisper_model
40
46
  self.whisper_language = self._normalize_whisper_language(whisper_language)
41
47
  self.utmos_path = utmos_model_path
42
48
  self.utmos_ckpt = utmos_ckpt_path
43
-
44
- self.whisper_model = None
45
- self.utmos_model = None
46
-
47
- if self.use_wer:
48
- self.wer_transform = jiwer.Compose([
49
- jiwer.ToLowerCase(),
50
- jiwer.RemovePunctuation(),
51
- jiwer.RemoveMultipleSpaces(),
52
- jiwer.Strip(),
53
- jiwer.RemoveEmptyStrings(),
54
- ])
55
-
49
+ self.zh_converter = OpenCC("t2s") if OpenCC is not None else None
50
+ self._warned_missing_opencc = False
51
+
52
+ self.whisper_model = None
53
+ self.utmos_model = None
54
+
55
+ if self.use_wer:
56
+ self.wer_transform = jiwer.Compose([
57
+ jiwer.ToLowerCase(),
58
+ jiwer.RemovePunctuation(),
59
+ jiwer.RemoveMultipleSpaces(),
60
+ jiwer.Strip(),
61
+ jiwer.RemoveEmptyStrings(),
62
+ ])
63
+
56
64
  @staticmethod
57
65
  def _normalize_whisper_language(language: Optional[str]) -> Optional[str]:
58
- if language is None:
59
- return None
60
- normalized = str(language).strip().lower()
61
- if normalized in {"zh", "zh-cn", "cmn", "<|cmn|>"}:
62
- return "zh"
63
- if normalized in {"en", "eng", "<|eng|>"}:
64
- return "en"
65
- if normalized in {"ja", "jp", "jpn"}:
66
- return "ja"
67
- if normalized in {"ko", "kor"}:
68
- return "ko"
69
- return normalized or None
70
-
71
- def _load_whisper(self):
72
- if self.whisper_model is None and self.use_wer:
73
- if not whisper:
74
- print("⚠️ Whisper is not installed, skipping load.")
75
- self.use_wer = False
76
- return
77
- print(f"⏳ Loading Whisper ({self.whisper_model_name})...")
78
- model_source, source_kind = resolve_pretrained_source(
79
- self.whisper_model_name,
80
- fallback_source=self.DEFAULT_WHISPER_MODEL,
81
- )
82
- print(f"Loading Whisper ({source_kind}) from {model_source}...")
83
- self.whisper_model = whisper.load_model(model_source, device=self.device)
84
-
85
- def _load_utmos(self):
86
- if self.utmos_model is None and self.use_utmos:
87
- print("⏳ Loading UTMOS model...")
88
- try:
89
- source = "github"
90
- repo_or_dir = "tarepan/SpeechMOS"
91
- if self.utmos_path and os.path.exists(self.utmos_path):
92
- source = "local"
93
- repo_or_dir = self.utmos_path
94
-
95
- load_pretrained = (self.utmos_ckpt is None)
96
- self.utmos_model = torch.hub.load(
97
- repo_or_dir, "utmos22_strong", source=source, trust_repo=True, pretrained=load_pretrained
98
- )
99
-
100
- if self.utmos_ckpt and os.path.exists(self.utmos_ckpt):
101
- state_dict = torch.load(self.utmos_ckpt, map_location=self.device)
102
- self.utmos_model.load_state_dict(state_dict)
103
-
104
- self.utmos_model.to(self.device).eval()
105
- except Exception as e:
106
- print(f"⚠️ UTMOS model failed to load: {e}")
107
- self.use_utmos = False
108
-
109
- def _transcribe(self, audio_paths: List[str]) -> List[str]:
110
- self._load_whisper()
111
- if not self.whisper_model:
112
- return [""] * len(audio_paths)
113
-
114
- transcribe_kwargs = {
115
- "fp16": torch.cuda.is_available() and "cuda" in self.device,
116
- "task": "transcribe",
117
- }
118
- if self.whisper_language:
119
- transcribe_kwargs["language"] = self.whisper_language
120
-
121
- results = []
122
- for path in tqdm(audio_paths, desc="Whisper transcription"):
123
- res = self.whisper_model.transcribe(path, **transcribe_kwargs)
124
- results.append(res["text"].strip())
125
- return results
126
-
127
- def _compute_utmos(self, audio_paths: List[str]) -> float:
128
- self._load_utmos()
129
- if not self.utmos_model: return 0.0
130
-
131
- scores = []
132
- target_sr = 16000
133
- for path in tqdm(audio_paths, desc="🎧 Computing UTMOS"):
134
- try:
135
- wave, sr = torchaudio.load(path)
136
- if sr != target_sr:
137
- wave = torchaudio.functional.resample(wave, sr, target_sr)
138
- if wave.shape[0] > 1:
139
- wave = torch.mean(wave, dim=0, keepdim=True)
140
-
141
- wave = wave.to(self.device)
142
- with torch.no_grad():
143
- s = self.utmos_model(wave, target_sr)
144
- scores.append(s.item())
145
- except Exception as e:
146
- print(f"⚠️ UTMOS Error on {path}: {e}")
147
-
148
- return sum(scores) / len(scores) if scores else 0.0
149
-
66
+ return whisper_language_code(language)
67
+
68
+ def _load_whisper(self):
69
+ if self.whisper_model is None and self.use_wer:
70
+ if not whisper:
71
+ print("⚠️ Whisper is not installed, skipping load.")
72
+ self.use_wer = False
73
+ return
74
+ print(f"⏳ Loading Whisper ({self.whisper_model_name})...")
75
+ model_source, source_kind = resolve_pretrained_source(
76
+ self.whisper_model_name,
77
+ fallback_source=self.DEFAULT_WHISPER_MODEL,
78
+ )
79
+ print(f"Loading Whisper ({source_kind}) from {model_source}...")
80
+ self.whisper_model = whisper.load_model(model_source, device=self.device)
81
+
82
+ def _load_utmos(self):
83
+ if self.utmos_model is None and self.use_utmos:
84
+ print("⏳ Loading UTMOS model...")
85
+ try:
86
+ source = "github"
87
+ repo_or_dir = "tarepan/SpeechMOS"
88
+ if self.utmos_path and os.path.exists(self.utmos_path):
89
+ source = "local"
90
+ repo_or_dir = self.utmos_path
91
+
92
+ load_pretrained = (self.utmos_ckpt is None)
93
+ self.utmos_model = torch.hub.load(
94
+ repo_or_dir, "utmos22_strong", source=source, trust_repo=True, pretrained=load_pretrained
95
+ )
96
+
97
+ if self.utmos_ckpt and os.path.exists(self.utmos_ckpt):
98
+ state_dict = torch.load(self.utmos_ckpt, map_location=self.device)
99
+ self.utmos_model.load_state_dict(state_dict)
100
+
101
+ self.utmos_model.to(self.device).eval()
102
+ except Exception as e:
103
+ print(f"⚠️ UTMOS model failed to load: {e}")
104
+ self.use_utmos = False
105
+
106
+ def _transcribe(self, audio_paths: List[str]) -> List[str]:
107
+ self._load_whisper()
108
+ if not self.whisper_model:
109
+ return [""] * len(audio_paths)
110
+
111
+ transcribe_kwargs = {
112
+ "fp16": torch.cuda.is_available() and "cuda" in self.device,
113
+ "task": "transcribe",
114
+ }
115
+ if self.whisper_language:
116
+ transcribe_kwargs["language"] = self.whisper_language
117
+
118
+ results = []
119
+ for path in tqdm(audio_paths, desc="Whisper transcription"):
120
+ res = self.whisper_model.transcribe(path, **transcribe_kwargs)
121
+ results.append(res["text"].strip())
122
+ return results
123
+
124
+ def _compute_utmos(self, audio_paths: List[str]) -> float:
125
+ self._load_utmos()
126
+ if not self.utmos_model: return 0.0
127
+
128
+ scores = []
129
+ target_sr = 16000
130
+ for path in tqdm(audio_paths, desc="🎧 Computing UTMOS"):
131
+ try:
132
+ wave, sr = torchaudio.load(path)
133
+ if sr != target_sr:
134
+ wave = torchaudio.functional.resample(wave, sr, target_sr)
135
+ if wave.shape[0] > 1:
136
+ wave = torch.mean(wave, dim=0, keepdim=True)
137
+
138
+ wave = wave.to(self.device)
139
+ with torch.no_grad():
140
+ s = self.utmos_model(wave, target_sr)
141
+ scores.append(s.item())
142
+ except Exception as e:
143
+ print(f"⚠️ UTMOS Error on {path}: {e}")
144
+
145
+ return sum(scores) / len(scores) if scores else 0.0
146
+
150
147
  def _preprocess_for_wer(self, text: str, lang: str) -> str:
148
+ normalized_lang = normalize_language_code(lang)
151
149
  text = self.wer_transform(text)
152
- if lang in ["zh", "ja", "ko"]:
150
+ if normalized_lang == "zh":
151
+ if self.zh_converter is not None:
152
+ text = self.zh_converter.convert(text)
153
+ elif not self._warned_missing_opencc:
154
+ print("Warning: OpenCC is not installed; Chinese CER will not normalize traditional/simplified variants.")
155
+ self._warned_missing_opencc = True
156
+ if speech_consistency_unit(normalized_lang) == "cer":
153
157
  text = text.replace(" ", "")
154
158
  return " ".join(list(text))
155
159
  return text
156
-
157
- def evaluate_all(self,
158
- target_audio: Union[List[str], str],
159
- target_text: Optional[Union[List[str], str]] = None,
160
- target_lang: str = "en") -> Dict[str, float]:
161
- """
162
- :param target_audio: List of paths or directory to the model-generated audio.
163
- :param target_text: Text generated synchronously by the model, used as reference for WER/CER calculation.
164
- :param target_lang: Target language for both audio and text.
165
- """
166
- results = {}
167
- print(f"\n--- Starting Speech Quality Evaluation (Target Lang: {target_lang}) ---")
168
-
169
- # Load target audio
170
- if isinstance(target_audio, str) and os.path.exists(target_audio) and os.path.isdir(target_audio):
171
- audio_paths = load_audio_from_folder(target_audio)
172
- else:
173
- audio_paths = target_audio if isinstance(target_audio, list) else [target_audio]
174
-
175
- # 1. evaluate UTMOS
176
- if self.use_utmos:
177
- print(" ➤ Computing UTMOS (Audio Naturalness)...")
178
- results["UTMOS"] = round(self._compute_utmos(audio_paths), 4)
179
-
180
- # 2. evaluate Consistency WER/CER
181
- if self.use_wer:
182
- if not target_text:
183
- print(" ⚠️ target_text (synchronously generated text) not provided. Text-speech consistency cannot be evaluated, skipping WER/CER.")
184
- else:
185
- texts = load_text_from_file_or_list(target_text, "target_text")
186
- if len(texts) != len(audio_paths):
187
- raise ValueError(f"Number of generated text lines ({len(texts)}) does not match number of audio files ({len(audio_paths)})!")
188
-
189
- print(" ➤ Transcribing audio, preparing consistency calculation...")
190
- asr_texts = self._transcribe(audio_paths)
191
-
192
- clean_refs = [self._preprocess_for_wer(t, target_lang) for t in texts]
193
- clean_hyps = [self._preprocess_for_wer(t, target_lang) for t in asr_texts]
194
-
160
+
161
+ def evaluate_all(self,
162
+ target_audio: Union[List[str], str],
163
+ target_text: Optional[Union[List[str], str]] = None,
164
+ target_lang: str = "en") -> Dict[str, float]:
165
+ """
166
+ :param target_audio: List of paths or directory to the model-generated audio.
167
+ :param target_text: Text generated synchronously by the model, used as reference for WER/CER calculation.
168
+ :param target_lang: Target language for both audio and text.
169
+ """
170
+ results = {}
171
+ print(f"\n--- Starting Speech Quality Evaluation (Target Lang: {target_lang}) ---")
172
+
173
+ # Load target audio
174
+ if isinstance(target_audio, str) and os.path.exists(target_audio) and os.path.isdir(target_audio):
175
+ audio_paths = load_audio_from_folder(target_audio)
176
+ else:
177
+ audio_paths = target_audio if isinstance(target_audio, list) else [target_audio]
178
+
179
+ # 1. evaluate UTMOS
180
+ if self.use_utmos:
181
+ print(" ➤ Computing UTMOS (Audio Naturalness)...")
182
+ results["UTMOS"] = round(self._compute_utmos(audio_paths), 4)
183
+
184
+ # 2. evaluate Consistency WER/CER
185
+ if self.use_wer:
186
+ if not target_text:
187
+ print(" ⚠️ target_text (synchronously generated text) not provided. Text-speech consistency cannot be evaluated, skipping WER/CER.")
188
+ else:
189
+ texts = load_text_from_file_or_list(target_text, "target_text")
190
+ if len(texts) != len(audio_paths):
191
+ raise ValueError(f"Number of generated text lines ({len(texts)}) does not match number of audio files ({len(audio_paths)})!")
192
+
193
+ print(" ➤ Transcribing audio, preparing consistency calculation...")
194
+ asr_texts = self._transcribe(audio_paths)
195
+
196
+ clean_refs = [self._preprocess_for_wer(t, target_lang) for t in texts]
197
+ clean_hyps = [self._preprocess_for_wer(t, target_lang) for t in asr_texts]
198
+
195
199
  if clean_refs:
196
200
  error_rate = jiwer.wer(clean_refs, clean_hyps)
197
- metric_name = "CER_Consistency" if target_lang in ['zh', 'ja', 'ko'] else "WER_Consistency"
201
+ metric_name = "CER_Consistency" if speech_consistency_unit(target_lang) == "cer" else "WER_Consistency"
198
202
  results[metric_name] = round(error_rate, 4)
199
- else:
200
- print(" ⚠️ The provided generated text is empty after cleaning, cannot compute distribution error rate.")
201
-
202
- return results
203
+ else:
204
+ print(" ⚠️ The provided generated text is empty after cleaning, cannot compute distribution error rate.")
205
+
206
+ return results
@@ -84,8 +84,9 @@ def load_audio_from_folder(folder_path: str, extensions: tuple = (".wav", ".mp3"
84
84
 
85
85
  # ==================== Evaluator Core Class ====================
86
86
 
87
- DEFAULT_COMET_MODEL = "Unbabel/wmt22-comet-da"
87
+ DEFAULT_COMET_MODEL = "Unbabel/wmt22-cometkiwi-da"
88
88
  DEFAULT_BLEURT_MODEL = "lucadiliello/BLEURT-20"
89
+ COMET_QE_METRIC_NAME = "COMETKiwi"
89
90
 
90
91
  class TranslationEvaluator:
91
92
  """
@@ -237,7 +238,7 @@ class TranslationEvaluator:
237
238
  Args:
238
239
  reference: Ground truth benchmark translation reference from the dataset (List or file path).
239
240
  target_text: Direct text translation output from the translation model (List or file path).
240
- source: Optional original source text (mandatory for COMET).
241
+ source: Optional original source text (mandatory for COMETKiwi/QE).
241
242
  target_lang: Target language, affects BLEU Tokenizer selection.
242
243
  """
243
244
  results = {}
@@ -276,13 +277,19 @@ class TranslationEvaluator:
276
277
  except Exception as e:
277
278
  results["BLEURT"] = -1.0
278
279
 
279
- # 4. COMET
280
- if self.use_comet and self.comet and final_src:
280
+ # 4. COMETKiwi reference-free QE. References may still be passed for
281
+ # BLEU/chrF++/BLEURT, but COMETKiwi must only receive source and MT.
282
+ if self.use_comet and self.comet:
283
+ if not final_src:
284
+ results[COMET_QE_METRIC_NAME] = -1.0
285
+ return {k: round(v, 4) if v >= 0 else v for k, v in results.items()}
286
+ if len(final_src) != len(final_text):
287
+ raise ValueError("Source and Target Text have different lengths")
281
288
  try:
282
- data = [{"src": s, "mt": t, "ref": r} for s, t, r in zip(final_src, final_text, final_ref)]
289
+ data = [{"src": s, "mt": t} for s, t in zip(final_src, final_text)]
283
290
  gpus = 1 if self.device.startswith("cuda") else 0
284
- results["COMET"] = self.comet.predict(data, batch_size=8, gpus=gpus).system_score
291
+ results[COMET_QE_METRIC_NAME] = self.comet.predict(data, batch_size=8, gpus=gpus).system_score
285
292
  except Exception as e:
286
- results["COMET"] = -1.0
293
+ results[COMET_QE_METRIC_NAME] = -1.0
287
294
 
288
295
  return {k: round(v, 4) if v >= 0 else v for k, v in results.items()}
@@ -1,14 +0,0 @@
1
- # This file contains code adapted from SimulEval:
2
- # https://github.com/facebookresearch/SimulEval
3
- #
4
- # SimulEval is licensed under the Creative Commons Attribution-ShareAlike
5
- # 4.0 International License (CC BY-SA 4.0):
6
- # https://creativecommons.org/licenses/by-sa/4.0/
7
- #
8
- # The adapted portions in this file are distributed under CC BY-SA 4.0.
9
- # Modifications were made for OpenSTBench.
10
-
11
- from .agent import GenericAgent, AgentPipeline
12
- from .basics import ReadAction, WriteAction, TextSegment, SpeechSegment
13
- from .cli import LatencyEvaluator
14
- from .metrics import register # Allows users to register custom Scorers
File without changes
File without changes
File without changes