Strict prompt construction in load_test.py (minimal)

llm_bench/load_test.py

@@ -26,6 +26,8 @@
 import gevent
 from locust.util.timespan import parse_timespan as _locust_parse_timespan
 
+from prefill_load_test import build_ids_to_length, build_pair_ids, load_chunks, split_chat_template
+
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
@@ -185,43 +187,85 @@ def __init__(
         chat: bool,
         num_tokens: int,
         common_tokens: int,
+        *,
+        strict: bool = True,
+        seed: int = 0,
     ):
         self._tokenizer = tokenizer
         self._tokenizer_path = tokenizer_path
         self._num_tokens = num_tokens
+        self._strict = strict
 
         self._all_limericks = []
         with open(path, "r") as f:
             text = f.read()
             lims = text.split("\n\n")
             for i, lim in enumerate(lims):
-                num_tokens = len(self._tokenizer.encode(lim, add_special_tokens=False))
-                self._all_limericks.append((lim, num_tokens))
+                lim_tokens = len(self._tokenizer.encode(lim, add_special_tokens=False))
+                self._all_limericks.append((lim, lim_tokens))
 
-        self._prefix = ""
         self._suffix = prompt
-        self._prefix_suffix_tokens = len(self._tokenizer.encode(prompt, add_special_tokens=False))
-        # Use deterministic selection (sequential iteration) to ensure all workers
-        # get the same prefix for the same common_tokens value
-        idx = 0
-        while self._prefix_suffix_tokens < common_tokens:
-            lim, num_tokens = self._all_limericks[idx % len(self._all_limericks)]
-            self._prefix += lim + "\n\n"
-            self._prefix_suffix_tokens += num_tokens
-            idx += 1
-
-        if chat:
-            empty_template_tokens = empty_chat_template_token_ids(self._tokenizer, self._tokenizer_path)
-            self._prefix_suffix_tokens += len(empty_template_tokens)
+
+        if strict:
+            self._rng = random.Random(seed)
+            dataset_name = "limericks" if path.endswith("limericks.txt") else "code"
+            self._chunks = load_chunks(dataset_name)
+            self._suffix_ids = self._tokenizer.encode(prompt, add_special_tokens=False)
+            content_budget = num_tokens - len(self._suffix_ids)
+            self._cached_tokens = min(common_tokens, content_budget)
+            self._base_ids = build_ids_to_length(tokenizer, self._chunks, content_budget)
+            if chat:
+                model_type = resolve_model_type(tokenizer_path)
+                self._chat_prefix, self._chat_suffix = split_chat_template(
+                    tokenizer, tokenizer_path, model_type
+                )
+            else:
+                self._chat_prefix, self._chat_suffix = [], []
+        else:
+            self._prefix = ""
+            self._prefix_suffix_tokens = len(self._tokenizer.encode(prompt, add_special_tokens=False))
+            # Use deterministic selection (sequential iteration) to ensure all workers
+            # get the same prefix for the same common_tokens value
+            idx = 0
+            while self._prefix_suffix_tokens < common_tokens:
+                lim, lim_tokens = self._all_limericks[idx % len(self._all_limericks)]
+                self._prefix += lim + "\n\n"
+                self._prefix_suffix_tokens += lim_tokens
+                idx += 1
+
+            if chat:
+                empty_template_tokens = empty_chat_template_token_ids(self._tokenizer, self._tokenizer_path)
+                self._prefix_suffix_tokens += len(empty_template_tokens)
 
     def __next__(self):
+        if self._strict:
+            content_tokens = self._num_tokens - len(self._suffix_ids)
+            ids = build_pair_ids(
+                self._chat_prefix,
+                self._chat_suffix,
+                self._base_ids,
+                self._tokenizer,
+                self._chunks,
+                content_tokens,
+                self._cached_tokens,
+                self._rng,
+            )
+            if self._suffix_ids:
+                chat_suffix_len = len(self._chat_suffix)
+                if chat_suffix_len:
+                    ids = ids[:-chat_suffix_len] + self._suffix_ids + ids[-chat_suffix_len:]
+                else:
+                    ids = ids + self._suffix_ids
+            assert len(ids) == self._num_tokens
+            return ids, self._num_tokens
+
         prompt_tokens = self._prefix_suffix_tokens
         prompt = self._prefix
         while prompt_tokens < self._num_tokens:
-            lim, num_tokens = self._all_limericks[random.randint(0, len(self._all_limericks) - 1)]
+            lim, lim_tokens = self._all_limericks[random.randint(0, len(self._all_limericks) - 1)]
 
             prompt += lim + "\n\n"
-            prompt_tokens += num_tokens
+            prompt_tokens += lim_tokens
         prompt += self._suffix
 
         return prompt, prompt_tokens
@@ -305,6 +349,7 @@ def _create_dataset(cls, options: argparse.Namespace):
                 chat=options.chat and not getattr(options, "rerank", False),
                 num_tokens=options.prompt_tokens,
                 common_tokens=common_tokens,
+                strict=not getattr(options, "rerank", False),
             )
         else:
             raise ValueError(f"Unknown dataset: {options.dataset}")
@@ -805,7 +850,9 @@ def format_payload(self, prompt, max_tokens, images):
             data["logprobs"] = self.parsed_options.logprobs
         if self.parsed_options.reasoning_effort is not None:
             data["reasoning_effort"] = self.parsed_options.reasoning_effort
-        if isinstance(prompt, str):
+        if isinstance(prompt, list) and prompt and isinstance(prompt[0], int):
+            data["prompt"] = prompt
+        elif isinstance(prompt, str):
             if self.parsed_options.chat:
                 if images is None:
                     data["messages"] = [{"role": "user", "content": prompt}]
@@ -943,12 +990,13 @@ def __init__(self, model, parsed_options):
             self._forced_generation_pool = itertools.cycle(self._load_forced_generation_texts(forced_gen_path))
         elif forced_gen_from_dataset:
             assert parsed_options.tokenizer is not None, "--tokenizer is required for --forced-generation-from-dataset"
+            tokenizer = InitTracker.load_tokenizer(parsed_options.tokenizer)
             ds = DatasetHolder.get_forced_generation_instance(
                 dataset=parsed_options.dataset,
                 tokenizer=parsed_options.tokenizer,
                 max_tokens=parsed_options.max_tokens,
             )
-            self._forced_generation_pool = (text for text, _tokens in ds)
+            self._forced_generation_pool = (tokenizer.decode(ids) for ids, _tokens in ds)
         else:
             self._forced_generation_pool = None
 
@@ -1409,7 +1457,11 @@ def _on_start(self):
         self.dataset = iter(dataset)
 
         tokenizer = InitTracker.load_tokenizer(self.environment.parsed_options.tokenizer)
-        self.prompt_tokenizer_tokens = len(tokenizer.encode(self._get_input()[0]))
+        sample_prompt = self._get_input()[0]
+        if isinstance(sample_prompt, list):
+            self.prompt_tokenizer_tokens = len(sample_prompt)
+        else:
+            self.prompt_tokenizer_tokens = len(tokenizer.encode(sample_prompt))
 
         # Override dataset with synthetic rerank documents if num_documents or tokens_per_document is set
         if self.environment.parsed_options.rerank and (
@@ -1516,8 +1568,12 @@ def _do_generate_text(self):
             print("---")
         t_start = time.perf_counter()
 
+        url = self.provider_formatter.get_url()
+        if isinstance(prompt, list) and prompt and isinstance(prompt[0], int):
+            url = "/v1/completions"
+
         with self.client.post(
-            self.provider_formatter.get_url(),
+            url,
             data=json.dumps(data),
             stream=True,
             catch_response=True,

llm_bench/test_load_test_strict.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Unit tests for strict token-id prompt construction in load_test.py."""
+
+import os
+import random
+import unittest
+from unittest import mock
+
+import transformers
+
+from load_test import OpenAIProvider, TranslationDataset
+from prefill_load_test import load_chunks
+
+
+class StrictPromptTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.tokenizer = transformers.AutoTokenizer.from_pretrained("gpt2")
+        cls.tokenizer.add_bos_token = False
+        cls.tokenizer.add_eos_token = False
+        cls.dataset_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "limericks.txt")
+
+    def _make_dataset(self, **kwargs):
+        defaults = dict(
+            path=self.dataset_path,
+            prompt="\n\nTranslate the limericks above to Spanish.",
+            tokenizer=self.tokenizer,
+            tokenizer_path="gpt2",
+            chat=False,
+            num_tokens=512,
+            common_tokens=0,
+            strict=True,
+            seed=1,
+        )
+        defaults.update(kwargs)
+        return TranslationDataset(**defaults)
+
+    def test_exact_prompt_length_no_chat(self):
+        ds = self._make_dataset()
+        ids, reported = next(iter(ds))
+        self.assertEqual(reported, 512)
+        self.assertEqual(len(ids), 512)
+
+    def test_exact_prompt_length_with_chat(self):
+        fake_prefix, fake_suffix = [1, 2, 3], [4, 5]
+        with mock.patch("load_test.resolve_model_type", return_value="gpt2"), mock.patch(
+            "load_test.split_chat_template", return_value=(fake_prefix, fake_suffix)
+        ):
+            ds = self._make_dataset(chat=True, seed=2)
+            ids, reported = next(iter(ds))
+        self.assertEqual(len(ids), 512)
+        self.assertEqual(ids[:3], fake_prefix)
+        self.assertEqual(ids[-2:], fake_suffix)
+
+    def test_shared_prefix_matches_cached_tokens(self):
+        ds = self._make_dataset(num_tokens=1000, common_tokens=600, seed=3)
+        first, _ = next(iter(ds))
+        second, _ = next(iter(ds))
+        self.assertEqual(first[:600], second[:600])
+        self.assertNotEqual(first, second)
+
+    def test_custom_prompt_suffix_is_included(self):
+        custom = "\n\nMy custom instruction."
+        ds = self._make_dataset(prompt=custom, num_tokens=256, common_tokens=0, seed=4)
+        suffix_ids = self.tokenizer.encode(custom, add_special_tokens=False)
+        ids, _ = next(iter(ds))
+        self.assertEqual(ids[-len(suffix_ids) :], suffix_ids)
+
+    def test_rerank_mode_returns_text_not_token_ids(self):
+        ds = self._make_dataset(strict=False, num_tokens=200, common_tokens=0)
+        prompt, _ = next(iter(ds))
+        self.assertIsInstance(prompt, str)
+
+    def test_format_payload_accepts_token_ids(self):
+        opts = mock.Mock(
+            chat=True,
+            rerank=False,
+            embeddings=False,
+            stream=False,
+            temperature=1.0,
+            n=1,
+            top_k=None,
+            logprobs=None,
+            reasoning_effort=None,
+            clear_assistant=False,
+        )
+        provider = OpenAIProvider("test-model", opts)
+        payload = provider.format_payload([10, 20, 30], max_tokens=16, images=None)
+        self.assertEqual(payload["prompt"], [10, 20, 30])
+        self.assertNotIn("messages", payload)
+
+    def test_rerank_format_payload_splits_text_not_token_ids(self):
+        opts = mock.Mock(
+            rerank=True,
+            rerank_query=None,
+            rerank_top_n=None,
+            rerank_return_documents=True,
+        )
+        provider = OpenAIProvider("test-model", opts)
+        payload = provider.format_payload("doc one\n\ndoc two", max_tokens=16, images=None)
+        self.assertEqual(payload["documents"], ["doc one", "doc two"])
+
+
+if __name__ == "__main__":
+    unittest.main()