OpenSTBench 1.2.0__tar.gz → 1.2.3__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.2.0/src/OpenSTBench.egg-info → openstbench-1.2.3}/PKG-INFO +1 -1
  2. {openstbench-1.2.0 → openstbench-1.2.3}/pyproject.toml +1 -1
  3. {openstbench-1.2.0 → openstbench-1.2.3/src/OpenSTBench.egg-info}/PKG-INFO +1 -1
  4. {openstbench-1.2.0 → openstbench-1.2.3}/src/OpenSTBench.egg-info/SOURCES.txt +1 -0
  5. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/__init__.py +4 -2
  6. openstbench-1.2.3/src/openstbench/language_policy.py +91 -0
  7. openstbench-1.2.3/src/openstbench/latency/__init__.py +39 -0
  8. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/speech_quality_evaluator.py +173 -184
  9. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/translation_evaluator.py +141 -79
  10. openstbench-1.2.0/src/openstbench/latency/__init__.py +0 -14
  11. {openstbench-1.2.0 → openstbench-1.2.3}/LICENSE +0 -0
  12. {openstbench-1.2.0 → openstbench-1.2.3}/README.md +0 -0
  13. {openstbench-1.2.0 → openstbench-1.2.3}/setup.cfg +0 -0
  14. {openstbench-1.2.0 → openstbench-1.2.3}/src/OpenSTBench.egg-info/dependency_links.txt +0 -0
  15. {openstbench-1.2.0 → openstbench-1.2.3}/src/OpenSTBench.egg-info/requires.txt +0 -0
  16. {openstbench-1.2.0 → openstbench-1.2.3}/src/OpenSTBench.egg-info/top_level.txt +0 -0
  17. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/_model_loading.py +0 -0
  18. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/emotion_evaluator.py +0 -0
  19. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/latency/agent.py +0 -0
  20. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/latency/basics.py +0 -0
  21. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/latency/cli.py +0 -0
  22. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/latency/instance.py +0 -0
  23. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/latency/metrics.py +0 -0
  24. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/latency/utils.py +0 -0
  25. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/paralinguistic_evaluator.py +0 -0
  26. {openstbench-1.2.0 → openstbench-1.2.3}/src/openstbench/speaker_similarity_evaluator.py +0 -0
  27. {openstbench-1.2.0 → openstbench-1.2.3}/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.2.0
3
+ Version: 1.2.3
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
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "OpenSTBench"
7
- version = "1.2.0"
7
+ version = "1.2.3"
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"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: OpenSTBench
3
- Version: 1.2.0
3
+ Version: 1.2.3
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
@@ -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
@@ -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.2.0"
12
+ __version__ = "1.2.3"
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,15 +1,16 @@
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:
@@ -19,143 +20,132 @@ try:
19
20
  from opencc import OpenCC
20
21
  except ImportError:
21
22
  OpenCC = None
22
-
23
- class SpeechQualityEvaluator:
24
- DEFAULT_WHISPER_MODEL = "medium"
25
- """
26
- Speech Quality and Consistency Evaluator.
27
- Specifically designed for:
28
- 1. UTMOS (Audio naturalness / perceived quality).
29
- 2. WER/CER Consistency Calculation (compares the model's generated text with its generated audio).
30
- """
31
- def __init__(self,
32
- use_wer: bool = True,
33
- use_utmos: bool = True,
34
- whisper_model: str = DEFAULT_WHISPER_MODEL,
35
- whisper_language: Optional[str] = None,
36
- utmos_model_path: Optional[str] = None,
37
- utmos_ckpt_path: Optional[str] = None,
38
- device: Optional[str] = None):
39
-
40
- self.device = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
41
- self.use_wer = use_wer
42
- self.use_utmos = use_utmos
43
-
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
+
44
45
  self.whisper_model_name = whisper_model
45
46
  self.whisper_language = self._normalize_whisper_language(whisper_language)
46
47
  self.utmos_path = utmos_model_path
47
48
  self.utmos_ckpt = utmos_ckpt_path
48
49
  self.zh_converter = OpenCC("t2s") if OpenCC is not None else None
49
50
  self._warned_missing_opencc = False
50
-
51
- self.whisper_model = None
52
- self.utmos_model = None
53
-
54
- if self.use_wer:
55
- self.wer_transform = jiwer.Compose([
56
- jiwer.ToLowerCase(),
57
- jiwer.RemovePunctuation(),
58
- jiwer.RemoveMultipleSpaces(),
59
- jiwer.Strip(),
60
- jiwer.RemoveEmptyStrings(),
61
- ])
62
-
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
+
63
64
  @staticmethod
64
65
  def _normalize_whisper_language(language: Optional[str]) -> Optional[str]:
65
- if language is None:
66
- return None
67
- normalized = str(language).strip().lower()
68
- if normalized in {"zh", "zh-cn", "cmn", "<|cmn|>"}:
69
- return "zh"
70
- if normalized in {"en", "eng", "<|eng|>"}:
71
- return "en"
72
- if normalized in {"ja", "jp", "jpn"}:
73
- return "ja"
74
- if normalized in {"ko", "kor"}:
75
- return "ko"
76
- return normalized or None
77
-
78
- def _load_whisper(self):
79
- if self.whisper_model is None and self.use_wer:
80
- if not whisper:
81
- print("⚠️ Whisper is not installed, skipping load.")
82
- self.use_wer = False
83
- return
84
- print(f"⏳ Loading Whisper ({self.whisper_model_name})...")
85
- model_source, source_kind = resolve_pretrained_source(
86
- self.whisper_model_name,
87
- fallback_source=self.DEFAULT_WHISPER_MODEL,
88
- )
89
- print(f"Loading Whisper ({source_kind}) from {model_source}...")
90
- self.whisper_model = whisper.load_model(model_source, device=self.device)
91
-
92
- def _load_utmos(self):
93
- if self.utmos_model is None and self.use_utmos:
94
- print("⏳ Loading UTMOS model...")
95
- try:
96
- source = "github"
97
- repo_or_dir = "tarepan/SpeechMOS"
98
- if self.utmos_path and os.path.exists(self.utmos_path):
99
- source = "local"
100
- repo_or_dir = self.utmos_path
101
-
102
- load_pretrained = (self.utmos_ckpt is None)
103
- self.utmos_model = torch.hub.load(
104
- repo_or_dir, "utmos22_strong", source=source, trust_repo=True, pretrained=load_pretrained
105
- )
106
-
107
- if self.utmos_ckpt and os.path.exists(self.utmos_ckpt):
108
- state_dict = torch.load(self.utmos_ckpt, map_location=self.device)
109
- self.utmos_model.load_state_dict(state_dict)
110
-
111
- self.utmos_model.to(self.device).eval()
112
- except Exception as e:
113
- print(f"⚠️ UTMOS model failed to load: {e}")
114
- self.use_utmos = False
115
-
116
- def _transcribe(self, audio_paths: List[str]) -> List[str]:
117
- self._load_whisper()
118
- if not self.whisper_model:
119
- return [""] * len(audio_paths)
120
-
121
- transcribe_kwargs = {
122
- "fp16": torch.cuda.is_available() and "cuda" in self.device,
123
- "task": "transcribe",
124
- }
125
- if self.whisper_language:
126
- transcribe_kwargs["language"] = self.whisper_language
127
-
128
- results = []
129
- for path in tqdm(audio_paths, desc="Whisper transcription"):
130
- res = self.whisper_model.transcribe(path, **transcribe_kwargs)
131
- results.append(res["text"].strip())
132
- return results
133
-
134
- def _compute_utmos(self, audio_paths: List[str]) -> float:
135
- self._load_utmos()
136
- if not self.utmos_model: return 0.0
137
-
138
- scores = []
139
- target_sr = 16000
140
- for path in tqdm(audio_paths, desc="🎧 Computing UTMOS"):
141
- try:
142
- wave, sr = torchaudio.load(path)
143
- if sr != target_sr:
144
- wave = torchaudio.functional.resample(wave, sr, target_sr)
145
- if wave.shape[0] > 1:
146
- wave = torch.mean(wave, dim=0, keepdim=True)
147
-
148
- wave = wave.to(self.device)
149
- with torch.no_grad():
150
- s = self.utmos_model(wave, target_sr)
151
- scores.append(s.item())
152
- except Exception as e:
153
- print(f"⚠️ UTMOS Error on {path}: {e}")
154
-
155
- return sum(scores) / len(scores) if scores else 0.0
156
-
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
+
157
147
  def _preprocess_for_wer(self, text: str, lang: str) -> str:
158
- normalized_lang = self._normalize_whisper_language(lang) or str(lang).strip().lower()
148
+ normalized_lang = normalize_language_code(lang)
159
149
  text = self.wer_transform(text)
160
150
  if normalized_lang == "zh":
161
151
  if self.zh_converter is not None:
@@ -163,55 +153,54 @@ class SpeechQualityEvaluator:
163
153
  elif not self._warned_missing_opencc:
164
154
  print("Warning: OpenCC is not installed; Chinese CER will not normalize traditional/simplified variants.")
165
155
  self._warned_missing_opencc = True
166
- if normalized_lang in ["zh", "ja", "ko"]:
156
+ if speech_consistency_unit(normalized_lang) == "cer":
167
157
  text = text.replace(" ", "")
168
158
  return " ".join(list(text))
169
159
  return text
170
-
171
- def evaluate_all(self,
172
- target_audio: Union[List[str], str],
173
- target_text: Optional[Union[List[str], str]] = None,
174
- target_lang: str = "en") -> Dict[str, float]:
175
- """
176
- :param target_audio: List of paths or directory to the model-generated audio.
177
- :param target_text: Text generated synchronously by the model, used as reference for WER/CER calculation.
178
- :param target_lang: Target language for both audio and text.
179
- """
180
- results = {}
181
- print(f"\n--- Starting Speech Quality Evaluation (Target Lang: {target_lang}) ---")
182
-
183
- # Load target audio
184
- if isinstance(target_audio, str) and os.path.exists(target_audio) and os.path.isdir(target_audio):
185
- audio_paths = load_audio_from_folder(target_audio)
186
- else:
187
- audio_paths = target_audio if isinstance(target_audio, list) else [target_audio]
188
-
189
- # 1. evaluate UTMOS
190
- if self.use_utmos:
191
- print(" ➤ Computing UTMOS (Audio Naturalness)...")
192
- results["UTMOS"] = round(self._compute_utmos(audio_paths), 4)
193
-
194
- # 2. evaluate Consistency WER/CER
195
- if self.use_wer:
196
- if not target_text:
197
- print(" ⚠️ target_text (synchronously generated text) not provided. Text-speech consistency cannot be evaluated, skipping WER/CER.")
198
- else:
199
- texts = load_text_from_file_or_list(target_text, "target_text")
200
- if len(texts) != len(audio_paths):
201
- raise ValueError(f"Number of generated text lines ({len(texts)}) does not match number of audio files ({len(audio_paths)})!")
202
-
203
- print(" ➤ Transcribing audio, preparing consistency calculation...")
204
- asr_texts = self._transcribe(audio_paths)
205
-
206
- clean_refs = [self._preprocess_for_wer(t, target_lang) for t in texts]
207
- clean_hyps = [self._preprocess_for_wer(t, target_lang) for t in asr_texts]
208
-
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
+
209
199
  if clean_refs:
210
200
  error_rate = jiwer.wer(clean_refs, clean_hyps)
211
- normalized_target_lang = self._normalize_whisper_language(target_lang) or str(target_lang).strip().lower()
212
- metric_name = "CER_Consistency" if normalized_target_lang in ['zh', 'ja', 'ko'] else "WER_Consistency"
201
+ metric_name = "CER_Consistency" if speech_consistency_unit(target_lang) == "cer" else "WER_Consistency"
213
202
  results[metric_name] = round(error_rate, 4)
214
- else:
215
- print(" ⚠️ The provided generated text is empty after cleaning, cannot compute distribution error rate.")
216
-
217
- return results
203
+ else:
204
+ print(" ⚠️ The provided generated text is empty after cleaning, cannot compute distribution error rate.")
205
+
206
+ return results
@@ -1,24 +1,28 @@
1
- import os
2
1
  import gc
3
2
  import json
3
+ import os
4
+ from pathlib import Path
5
+ from typing import Dict, List, Optional, Union
6
+
4
7
  import numpy as np
5
8
  import sacrebleu
6
9
  import torch
7
- from typing import Dict, List, Optional, Union
8
- from pathlib import Path
9
10
 
10
11
  from ._model_loading import resolve_pretrained_source
11
12
 
13
+
12
14
  # ==================== Configuration ====================
13
15
 
14
16
  CACHE_PATHS = {
15
17
  "huggingface": os.path.expanduser("~/.cache/huggingface/hub"),
16
18
  }
17
19
 
20
+
18
21
  # ==================== Optional Dependencies ====================
19
22
 
20
23
  try:
21
24
  from bleurt_pytorch import BleurtConfig, BleurtForSequenceClassification, BleurtTokenizer
25
+
22
26
  HAS_BLEURT = True
23
27
  except ImportError:
24
28
  HAS_BLEURT = False
@@ -29,46 +33,50 @@ except ImportError:
29
33
  download_model = None
30
34
  load_from_checkpoint = None
31
35
 
36
+
32
37
  # ==================== Input Loaders ====================
33
38
 
34
39
  def load_text_from_file_or_list(input_data: Union[str, List[str]], name: str = "text") -> List[str]:
35
40
 
36
41
  if isinstance(input_data, list):
37
42
  return input_data
38
-
43
+
39
44
  if not isinstance(input_data, str):
40
45
  raise ValueError(f"{name} must be a file path (str) or a list of strings (List[str])")
41
46
 
42
47
  path = Path(input_data)
43
48
  if not path.exists():
44
49
  raise FileNotFoundError(f"{name} file not found: {input_data}")
45
-
50
+
46
51
  suffix = path.suffix.lower()
47
-
52
+
48
53
  if suffix == ".json":
49
54
  with open(path, "r", encoding="utf-8") as f:
50
55
  data = json.load(f)
51
-
56
+
52
57
  if isinstance(data, list):
53
- if len(data) == 0: return []
54
- if isinstance(data[0], str): return data
58
+ if len(data) == 0:
59
+ return []
60
+ if isinstance(data[0], str):
61
+ return data
55
62
  if isinstance(data[0], dict):
56
63
  for key in ["target_text", "hypothesis", "text", "ref", "reference", "src", "source"]:
57
64
  if key in data[0]:
58
65
  return [item[key] for item in data]
59
- raise ValueError(f"JSON list items do not contain common text fields")
66
+ raise ValueError("JSON list items do not contain common text fields")
60
67
 
61
68
  if isinstance(data, dict):
62
69
  for key in ["target_text", "hypothesis", "text", "ref", "reference", "src", "source"]:
63
70
  if key in data:
64
71
  return data[key]
65
72
  raise ValueError("JSON dictionary does not contain common text fields")
66
-
73
+
67
74
  raise ValueError("Unsupported JSON format")
68
75
  else:
69
76
  with open(path, "r", encoding="utf-8") as f:
70
77
  return [line.strip() for line in f if line.strip()]
71
78
 
79
+
72
80
  def load_audio_from_folder(folder_path: str, extensions: tuple = (".wav", ".mp3", ".flac")) -> List[str]:
73
81
  folder = Path(folder_path)
74
82
  if not folder.exists():
@@ -85,23 +93,31 @@ def load_audio_from_folder(folder_path: str, extensions: tuple = (".wav", ".mp3"
85
93
  # ==================== Evaluator Core Class ====================
86
94
 
87
95
  DEFAULT_COMET_MODEL = "Unbabel/wmt22-comet-da"
96
+ DEFAULT_COMET_QE_MODEL = "Unbabel/wmt22-cometkiwi-da"
88
97
  DEFAULT_BLEURT_MODEL = "lucadiliello/BLEURT-20"
98
+ COMET_METRIC_NAME = "COMET"
99
+ COMET_QE_METRIC_NAME = "COMETKiwi"
100
+
89
101
 
90
102
  class TranslationEvaluator:
91
103
  """
92
- Text-side Translation Quality Evaluator: Supports Text Translation Extraction and Comparison (BLEU, COMET, BLEURT...)
104
+ Text-side Translation Quality Evaluator: Supports text translation metrics
105
+ such as BLEU, chrF++, COMET, COMETKiwi, and BLEURT.
93
106
  """
94
107
 
95
- def __init__(self,
96
- use_bleu: bool = True,
97
- use_chrf: bool = True,
98
- use_comet: bool = True,
99
- use_bleurt: bool = False,
100
- comet_model: str = DEFAULT_COMET_MODEL,
101
- bleurt_path: Optional[str] = None,
102
- bleurt_model: Optional[str] = None,
103
- device: Optional[str] = None):
104
-
108
+ def __init__(
109
+ self,
110
+ use_bleu: bool = True,
111
+ use_chrf: bool = True,
112
+ use_comet: bool = True,
113
+ use_bleurt: bool = False,
114
+ comet_model: str = DEFAULT_COMET_MODEL,
115
+ comet_qe_model: str = DEFAULT_COMET_QE_MODEL,
116
+ bleurt_path: Optional[str] = None,
117
+ bleurt_model: Optional[str] = None,
118
+ device: Optional[str] = None,
119
+ ):
120
+
105
121
  self.use_bleu = use_bleu
106
122
  self.use_chrf = use_chrf
107
123
  self.use_comet = use_comet
@@ -113,28 +129,30 @@ class TranslationEvaluator:
113
129
  self.device = "cuda"
114
130
  else:
115
131
  self.device = "cpu"
116
-
117
- print(f"🚀 Initializing Translation Evaluator (Text-only) on {self.device}...")
118
132
 
119
- # Load language models
133
+ print(f"Initializing Translation Evaluator (Text-only) on {self.device}...")
134
+
135
+ self.comet_model = comet_model
136
+ self.comet_qe_model = comet_qe_model
120
137
  self.comet = None
121
- if self.use_comet:
122
- self.comet = self._load_comet(comet_model)
138
+ self.comet_qe = None
123
139
 
124
140
  self.bleurt_model = None
125
141
  self.bleurt_tokenizer = None
126
142
  if self.use_bleurt:
127
143
  self._load_bleurt(bleurt_path, bleurt_model)
128
144
 
129
- print("Translation text evaluation metrics system is ready!")
145
+ print("Translation text evaluation metrics system is ready!")
146
+
147
+ def __enter__(self):
148
+ return self
130
149
 
131
- def __enter__(self): return self
132
- def __exit__(self, exc_type, exc_val, exc_tb):
150
+ def __exit__(self, exc_type, exc_val, exc_tb):
133
151
  self.cleanup()
134
152
  return False
135
153
 
136
154
  def cleanup(self):
137
- for attr in ['comet', 'bleurt_model', 'bleurt_tokenizer']:
155
+ for attr in ["comet", "comet_qe", "bleurt_model", "bleurt_tokenizer"]:
138
156
  if hasattr(self, attr) and getattr(self, attr) is not None:
139
157
  delattr(self, attr)
140
158
  if torch.cuda.is_available():
@@ -158,38 +176,38 @@ class TranslationEvaluator:
158
176
  return str(ckpt.resolve())
159
177
  return None
160
178
 
161
- def _load_comet(self, model_name: str):
179
+ def _load_comet(self, model_name: str, *, fallback_model: str, metric_name: str):
162
180
  if not download_model:
163
- print("⚠️ COMET is not installed, skipping...")
181
+ print("COMET is not installed, skipping...")
164
182
  return None
165
183
  try:
166
184
  model_source, source_kind = resolve_pretrained_source(
167
185
  model_name,
168
- fallback_source=DEFAULT_COMET_MODEL,
186
+ fallback_source=fallback_model,
169
187
  )
170
188
  local_ckpt = self._resolve_local_comet_checkpoint(model_source)
171
189
  if local_ckpt is not None:
172
- print(f"[Local] Loading COMET: {local_ckpt}")
190
+ print(f"[Local] Loading {metric_name}: {local_ckpt}")
173
191
  model = load_from_checkpoint(local_ckpt)
174
192
  else:
175
- remote_source = model_source if source_kind == "remote" else DEFAULT_COMET_MODEL
193
+ remote_source = model_source if source_kind == "remote" else fallback_model
176
194
  cache = os.path.join(CACHE_PATHS["huggingface"], f"models--{remote_source.replace('/', '--')}")
177
195
  status = "[Local]" if os.path.exists(cache) else "[Online]"
178
- print(f"{status} Loading COMET: {model_name}")
179
- print(f"Loading COMET ({status}) from {remote_source}")
196
+ print(f"{status} Loading {metric_name}: {model_name}")
197
+ print(f"Loading {metric_name} ({status}) from {remote_source}")
180
198
  model = load_from_checkpoint(download_model(remote_source))
181
199
  if self.device.startswith("cuda"):
182
200
  model = model.to(self.device)
183
201
  return model
184
202
  except Exception as e:
185
- print(f" COMET loading failed: {e}")
203
+ print(f"{metric_name} loading failed: {e}")
186
204
  return None
187
205
 
188
206
  def _load_bleurt(self, path: Optional[str], model_name: Optional[str]):
189
207
  if not HAS_BLEURT:
190
- print("⚠️ bleurt-pytorch is not installed, skipping...")
208
+ print("bleurt-pytorch is not installed, skipping...")
191
209
  return
192
-
210
+
193
211
  if path:
194
212
  model_source, _source_kind = resolve_pretrained_source(
195
213
  path,
@@ -197,30 +215,40 @@ class TranslationEvaluator:
197
215
  )
198
216
  else:
199
217
  model_source, _source_kind = resolve_pretrained_source(model_name or DEFAULT_BLEURT_MODEL)
200
- print(f"Loading BLEURT: {model_source}")
201
-
218
+ print(f"Loading BLEURT: {model_source}")
219
+
202
220
  try:
203
221
  self.bleurt_tokenizer = BleurtTokenizer.from_pretrained(model_source)
204
222
  self.bleurt_model = BleurtForSequenceClassification.from_pretrained(model_source)
205
223
  self.bleurt_model = self.bleurt_model.to(self.device).eval()
206
224
  except Exception as e:
207
- print(f"BLEURT loading failed: {e}")
208
-
225
+ print(f"BLEURT loading failed: {e}")
209
226
 
210
227
  def _get_bleu_tokenizer_name(self, lang: str) -> str:
211
- if lang == 'zh': return 'zh'
212
- elif lang == 'ja': return 'ja-mecab'
213
- elif lang == 'ko': return 'ko-mecab'
214
- else: return '13a'
228
+ if lang == "zh":
229
+ return "zh"
230
+ elif lang == "ja":
231
+ return "ja-mecab"
232
+ elif lang == "ko":
233
+ return "ko-mecab"
234
+ else:
235
+ return "13a"
215
236
 
216
237
  def _compute_bleurt_score(self, references: List[str], candidates: List[str]) -> float:
217
238
  all_scores = []
218
239
  batch_size = 32
219
240
  for i in range(0, len(references), batch_size):
220
- br = references[i:i+batch_size]
221
- bc = candidates[i:i+batch_size]
241
+ br = references[i : i + batch_size]
242
+ bc = candidates[i : i + batch_size]
222
243
  with torch.no_grad():
223
- inputs = self.bleurt_tokenizer(br, bc, padding='longest', truncation=True, max_length=512, return_tensors='pt')
244
+ inputs = self.bleurt_tokenizer(
245
+ br,
246
+ bc,
247
+ padding="longest",
248
+ truncation=True,
249
+ max_length=512,
250
+ return_tensors="pt",
251
+ )
224
252
  inputs = {k: v.to(self.device) for k, v in inputs.items()}
225
253
  scores = self.bleurt_model(**inputs).logits.flatten().tolist()
226
254
  all_scores.extend(scores)
@@ -228,61 +256,95 @@ class TranslationEvaluator:
228
256
 
229
257
  def evaluate_all(
230
258
  self,
231
- reference: Union[List[str], str],
232
- target_text: Union[List[str], str],
259
+ reference: Optional[Union[List[str], str]] = None,
260
+ target_text: Optional[Union[List[str], str]] = None,
233
261
  source: Optional[Union[List[str], str]] = None,
234
- target_lang: str = "en"
262
+ target_lang: str = "en",
235
263
  ) -> Dict[str, float]:
236
264
  """
237
265
  Args:
238
- reference: Ground truth benchmark translation reference from the dataset (List or file path).
239
- target_text: Direct text translation output from the translation model (List or file path).
240
- source: Optional original source text (mandatory for COMET).
241
- target_lang: Target language, affects BLEU Tokenizer selection.
266
+ reference: Optional benchmark translation reference from the dataset.
267
+ target_text: Direct text translation output from the translation model.
268
+ source: Optional original source text, mandatory for COMET and COMETKiwi.
269
+ target_lang: Target language, affects BLEU tokenizer selection.
242
270
  """
243
271
  results = {}
244
272
  print(f"\n--- Starting Text Translation Quality Evaluation (Target Lang: {target_lang}) ---")
245
273
 
246
- final_ref = load_text_from_file_or_list(reference, "Reference")
274
+ if target_text is None:
275
+ raise ValueError("target_text is required")
247
276
  final_text = load_text_from_file_or_list(target_text, "Target Text")
248
-
277
+
278
+ final_ref = None
279
+ if reference is not None:
280
+ loaded_ref = load_text_from_file_or_list(reference, "Reference")
281
+ final_ref = loaded_ref if loaded_ref else None
282
+
249
283
  final_src = None
250
284
  if source:
251
285
  final_src = load_text_from_file_or_list(source, "Source")
252
- if len(final_src) != len(final_ref):
253
- raise ValueError("Source and Reference have different lengths")
254
-
255
- if len(final_text) != len(final_ref):
286
+ if len(final_src) != len(final_text):
287
+ raise ValueError("Source and Target Text have different lengths")
288
+
289
+ if final_ref is not None and len(final_text) != len(final_ref):
256
290
  raise ValueError("Target Text and Reference have different lengths")
257
291
 
258
292
  # 1. sacreBLEU
259
- if self.use_bleu:
293
+ if self.use_bleu and final_ref is not None:
260
294
  tokenizer_name = self._get_bleu_tokenizer_name(target_lang)
261
295
  try:
262
296
  results["sacreBLEU"] = sacrebleu.corpus_bleu(final_text, [final_ref], tokenize=tokenizer_name).score
263
- except Exception as e:
297
+ except Exception:
264
298
  results["sacreBLEU"] = -1.0
265
299
 
266
300
  # 2. chrF++
267
- if self.use_chrf:
301
+ if self.use_chrf and final_ref is not None:
268
302
  try:
269
303
  results["chrF++"] = sacrebleu.corpus_chrf(final_text, [final_ref], word_order=2).score
270
- except: results["chrF++"] = -1.0
304
+ except Exception:
305
+ results["chrF++"] = -1.0
271
306
 
272
307
  # 3. BLEURT
273
- if self.use_bleurt and self.bleurt_model:
308
+ if self.use_bleurt and self.bleurt_model and final_ref is not None:
274
309
  try:
275
310
  results["BLEURT"] = self._compute_bleurt_score(final_ref, final_text)
276
- except Exception as e:
311
+ except Exception:
277
312
  results["BLEURT"] = -1.0
278
313
 
279
- # 4. COMET
280
- if self.use_comet and self.comet and final_src:
281
- try:
282
- data = [{"src": s, "mt": t, "ref": r} for s, t, r in zip(final_src, final_text, final_ref)]
314
+ # 4. COMET and COMETKiwi.
315
+ if self.use_comet:
316
+ if not final_src:
317
+ if final_ref is not None:
318
+ results[COMET_METRIC_NAME] = -1.0
319
+ results[COMET_QE_METRIC_NAME] = -1.0
320
+ else:
283
321
  gpus = 1 if self.device.startswith("cuda") else 0
284
- results["COMET"] = self.comet.predict(data, batch_size=8, gpus=gpus).system_score
285
- except Exception as e:
286
- results["COMET"] = -1.0
287
-
322
+
323
+ if final_ref is not None:
324
+ if self.comet is None:
325
+ self.comet = self._load_comet(
326
+ self.comet_model,
327
+ fallback_model=DEFAULT_COMET_MODEL,
328
+ metric_name=COMET_METRIC_NAME,
329
+ )
330
+ if self.comet:
331
+ try:
332
+ data = [{"src": s, "mt": t, "ref": r} for s, t, r in zip(final_src, final_text, final_ref)]
333
+ results[COMET_METRIC_NAME] = self.comet.predict(data, batch_size=8, gpus=gpus).system_score
334
+ except Exception:
335
+ results[COMET_METRIC_NAME] = -1.0
336
+
337
+ if self.comet_qe is None:
338
+ self.comet_qe = self._load_comet(
339
+ self.comet_qe_model,
340
+ fallback_model=DEFAULT_COMET_QE_MODEL,
341
+ metric_name=COMET_QE_METRIC_NAME,
342
+ )
343
+ if self.comet_qe:
344
+ try:
345
+ data = [{"src": s, "mt": t} for s, t in zip(final_src, final_text)]
346
+ results[COMET_QE_METRIC_NAME] = self.comet_qe.predict(data, batch_size=8, gpus=gpus).system_score
347
+ except Exception:
348
+ results[COMET_QE_METRIC_NAME] = -1.0
349
+
288
350
  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