fix: simulator returns zeros on transformers >= 5 and BOS-less tokenizers

delphi/clients/offline.py

@@ -113,15 +113,19 @@ async def process_func(
         statistics = []
 
         for batch in batches:
+            # return_dict=False: transformers 5 returns a BatchEncoding by default
             prompt = self.tokenizer.apply_chat_template(
-                batch, add_generation_prompt=True, tokenize=True
+                batch, add_generation_prompt=True, tokenize=True, return_dict=False
             )
             prompt = TokensPrompt(prompt_token_ids=prompt)
             prompts.append(prompt)
             if self.statistics:
                 non_cached_tokens = len(
                     self.tokenizer.apply_chat_template(
-                        batch[-1:], add_generation_prompt=True, tokenize=True  # type: ignore
+                        batch[-1:],  # type: ignore
+                        add_generation_prompt=True,
+                        tokenize=True,
+                        return_dict=False,
                     )
                 )
                 statistics.append(

delphi/scorers/simulator/simulation/simulator.py

@@ -123,11 +123,12 @@ async def _simulate_with_logprobs(self, tokens: List[str]):
         }
 
         response = await self.client.generate(prompt, **sampling_params)
+        # return_dict=False: transformers 5 returns a BatchEncoding by default
         tokenized_prompt = self.client.tokenizer.apply_chat_template(
-            prompt, add_generation_prompt=True
+            prompt, add_generation_prompt=True, return_dict=False
         )
-        # Use [1] to skip BOS token that tokenizer.encode() may prepend
-        tab_token = self.client.tokenizer.encode("\t")[1]
+        # Not every tokenizer prepends a BOS token, so encode without special tokens
+        tab_token = self.client.tokenizer.encode("\t", add_special_tokens=False)[-1]
 
         result = parse_simulation_response(
             response, tokenized_prompt, tab_token, tokens

tests/test_scorers/test_simulator_tokenization.py

@@ -0,0 +1,47 @@
+import asyncio
+import math
+from types import SimpleNamespace
+
+import pytest
+from transformers import AutoTokenizer
+
+from delphi.scorers.simulator.simulation.simulator import NeuronSimulator
+
+
+class _UniformLogprobClient:
+    """Returns prompt logprobs with P("7") = P("0") = 0.5 at every position."""
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    async def generate(self, prompt, **kwargs):
+        prompt_ids = self.tokenizer.apply_chat_template(
+            prompt, add_generation_prompt=True, return_dict=False
+        )
+        logprobs = {
+            0: SimpleNamespace(decoded_token="7", logprob=math.log(0.5)),
+            1: SimpleNamespace(decoded_token="0", logprob=math.log(0.5)),
+        }
+        return SimpleNamespace(prompt_logprobs=[logprobs] * (len(prompt_ids) + 1))
+
+
+@pytest.mark.parametrize(
+    "tokenizer_name",
+    [
+        # prepends a BOS token
+        "hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4",
+        # does not prepend a BOS token
+        "Qwen/Qwen2.5-0.5B-Instruct",
+    ],
+)
+def test_logprob_simulation_reads_every_token(tokenizer_name: str):
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    simulator = NeuronSimulator(
+        _UniformLogprobClient(tokenizer), "cats", use_logprobs=True
+    )
+    tokens = ["The", " cat", " sat"]
+
+    simulation = asyncio.run(simulator.simulate(tokens))
+
+    # A parsing failure would be swallowed by `simulate` and return all zeros
+    assert simulation.expected_activations == pytest.approx([3.5, 3.5, 3.5])