rebase fixes

Author: TJ5

genai_bench/cli/cli.py

@@ -10,11 +10,10 @@
 
 from genai_bench.analysis.excel_report import create_workbook
 from genai_bench.analysis.experiment_loader import load_one_experiment
+from genai_bench.analysis.flexible_plot_report import plot_experiment_data_flexible
 from genai_bench.analysis.plot_report import (
     plot_single_scenario_inference_speed_vs_throughput,
 )
-from genai_bench.analysis.flexible_plot_report import plot_experiment_data_flexible
-
 from genai_bench.auth.unified_factory import UnifiedAuthFactory
 from genai_bench.cli.option_groups import (
     api_options,

tests/sampling/test_text.py

@@ -267,12 +267,14 @@ def test_sample_chat_prefix_ratio_request(self):
 
         # Mock encode to return list with length equal to number of characters in input
         def mock_encode(text, add_special_tokens=False):
-            return [1] * len(text)
+            # ignore space
+            encoded_text = [1] * len(text.replace(" ", ""))
+            return encoded_text
 
         self.tokenizer.encode = mock_encode
 
         # Mock decode to return the original text
-        def mock_decode(tokens):
+        def mock_decode(tokens, skip_special_tokens=True):
             if isinstance(tokens, list):
                 return "a" * len(tokens)  # Return 'a' repeated for the token count
             return "decoded_text"
@@ -290,7 +292,6 @@ def mock_decode(tokens):
             model=self.model,
             output_modality=self.output_modality,
             data=self.test_data,
-            use_scenario=True,
             prompt_prefix_ratio=0.5,  # 50% of 20 tokens = 10 tokens
         )
         result = prefix_sampler.sample(scenario)
@@ -299,7 +300,7 @@ def mock_decode(tokens):
         self.assertTrue(isinstance(result.prompt, str))
         self.assertGreater(len(result.prompt), 0)
         self.assertTrue(result.prompt.startswith(prefix_sampler.prefix))
-        self.assertEqual(len(result.prompt), 20)
+        self.assertEqual(len(mock_encode(result.prompt)), 20)
 
     def test_short_prompt_request(self):
         """Test that short prompts are handled correctly."""
@@ -335,7 +336,6 @@ def test_empty_dataset(self):
             model=self.model,
             output_modality=self.output_modality,
             data=[],
-            use_scenario=True,
         )
         scenario = NormalDistribution(10, 0, 10, 0)