Support Fun-CosyVoice3 eagle3 speculative decoding

angelslim/compressor/speculative/benchmark/pytorch/__init__.py

@@ -12,6 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .benchmark_engine import BenchmarkConfig, BenchmarkEngine, BenchmarkMode
+from .benchmark_engine import (
+    BenchmarkConfig,
+    BenchmarkEngine,
+    BenchmarkMode,
+    TTSBenchmarkEngine,
+)
 
-__all__ = ["BenchmarkEngine", "BenchmarkConfig", "BenchmarkMode"]
+__all__ = ["BenchmarkEngine", "TTSBenchmarkEngine", "BenchmarkConfig", "BenchmarkMode"]

angelslim/compressor/speculative/benchmark/pytorch/benchmark_engine.py

@@ -25,7 +25,11 @@
 from angelslim.utils.lazy_imports import fastchat, ray
 
 from .generate_baseline_answer import get_model_answers as get_baseline_answers
+from .generate_baseline_answer import get_tts_answers as get_tts_baseline_answers
+from .generate_baseline_answer import get_tts_audios as get_tts_baseline_audios
 from .generate_eagle_answer import get_model_answers as get_eagle_answers
+from .generate_eagle_answer import get_tts_answers as get_tts_eagle_answers
+from .generate_eagle_answer import get_tts_audios as get_tts_eagle_audios
 
 
 class BenchmarkMode(Enum):
@@ -77,6 +81,10 @@ class BenchmarkConfig:
     # Batch settings
     batch_size: int = 1
 
+    # TTS settings
+    is_tts: bool = False
+    generate_audio: bool = False
+
 
 class BenchmarkEngine:
     """Core benchmark engine for speculative decoding evaluation"""
@@ -343,6 +351,10 @@ def _create_args_namespace(self, mode: str) -> argparse.Namespace:
 
         args.early_stop_method = self.config.early_stop_method
 
+        # TTS settings
+        args.is_tts = self.config.is_tts
+        args.generate_audio = self.config.generate_audio
+
         return args
 
     def _get_question_file_path(self) -> str:
@@ -397,3 +409,119 @@ def get_performance_summary(self) -> str:
         summary.append(f"Analysis Report: {self.analysis_file}")
 
         return "\n".join(summary)
+
+
+class TTSBenchmarkEngine(BenchmarkEngine):
+    """Core benchmark engine for speculative decoding evaluation"""
+
+    def _run_eagle_benchmark(self):
+        """Run Eagle speculative decoding benchmark"""
+        args = self._create_args_namespace("eagle")
+
+        questions = fastchat.llm_judge.common.load_questions(
+            self._get_question_file_path(),
+            self.config.question_begin,
+            self.config.question_end,
+        )
+
+        use_ray = self.config.num_gpus_total // self.config.num_gpus_per_model > 1
+        get_answers_func = (
+            ray.remote(num_gpus=self.config.num_gpus_per_model)(
+                get_tts_eagle_answers
+            ).remote
+            if use_ray
+            else get_tts_eagle_answers
+        )
+
+        chunk_size = len(questions) // (
+            self.config.num_gpus_total // self.config.num_gpus_per_model
+        )
+        ans_handles = [
+            get_answers_func(
+                f"{self.config.model_id}-temperature-{self.config.temperature}",
+                questions[i : i + chunk_size],
+                self.eagle_file,
+                self.config.num_choices,
+                self.config.temperature,
+                args,
+            )
+            for i in range(0, len(questions), chunk_size)
+        ]
+
+        if use_ray:
+            ray.get(ans_handles)
+
+        self._reorg_answer_file(self.eagle_file)
+        self.results["eagle_file"] = self.eagle_file
+
+        if self.config.generate_audio:
+            self._generate_audio("eagle")
+
+    def _run_baseline_benchmark(self):
+        """Run baseline benchmark"""
+        args = self._create_args_namespace("baseline")
+
+        questions = fastchat.llm_judge.common.load_questions(
+            self._get_question_file_path(),
+            self.config.question_begin,
+            self.config.question_end,
+        )
+
+        use_ray = self.config.num_gpus_total // self.config.num_gpus_per_model > 1
+        get_answers_func = (
+            ray.remote(num_gpus=self.config.num_gpus_per_model)(
+                get_tts_baseline_answers
+            ).remote
+            if use_ray
+            else get_tts_baseline_answers
+        )
+
+        chunk_size = len(questions) // (
+            self.config.num_gpus_total // self.config.num_gpus_per_model
+        )
+        ans_handles = [
+            get_answers_func(
+                f"{self.config.model_id}-temperature-{self.config.temperature}",
+                questions[i : i + chunk_size],
+                self.baseline_file,
+                self.config.num_choices,
+                self.config.temperature,
+                args,
+            )
+            for i in range(0, len(questions), chunk_size)
+        ]
+
+        if use_ray:
+            ray.get(ans_handles)
+
+        self._reorg_answer_file(self.baseline_file)
+        self.results["baseline_file"] = self.baseline_file
+
+        if self.config.generate_audio:
+            self._generate_audio("baseline")
+
+    def _calculate_metrics(self) -> Dict[str, Any]:
+        """Calculate acceptance length and speedup ratio"""
+        metrics = {}
+
+        # Calculate acceptance length from Eagle results
+        if os.path.exists(self.eagle_file):
+            metrics["acceptance_length"] = self._calculate_acceptance_length(
+                self.eagle_file
+            )
+
+        return metrics
+
+    def _generate_audio(self, mode):
+        args = self._create_args_namespace(mode)
+
+        answers = fastchat.llm_judge.common.load_questions(
+            args.answer_file,
+            self.config.question_begin,
+            self.config.question_end,
+        )
+
+        if mode == "baseline":
+            get_tts_baseline_audios(answers, args.answer_file, args)
+        else:
+            get_tts_eagle_audios(answers, args.answer_file, args)