interp-engine 1.7.2__tar.gz → 1.7.3__tar.gz

{interp_engine-1.7.2 → interp_engine-1.7.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: interp-engine
-Version: 1.7.2
+Version: 1.7.3
 Summary: A fast, standardized, and easy to use interpretability engine.
 Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
 Project-URL: Repository, https://github.com/decoderesearch/interp-engine

{interp_engine-1.7.2 → interp_engine-1.7.3}/docs/USAGE.md

@@ -215,7 +215,9 @@ one entry in `chat_formatters.CODE_CHAT_FORMATS`.
 
 To attribute *tokens* to messages there are two methods, and the difference matters. `message_spans`
 gives per-token role, channel and section (`header` / `content` / `footer`), leaving the trailing
-generation scaffold owned by no message — use it to read or display structure. `message_partition`
+generation scaffold owned by no message — use it to read or display structure. A system turn the
+template injects on its own (Llama's knowledge-cutoff preamble, Qwen2.5's default persona) is
+tagged `role="system"` with `message_index=None`, so it can be shown as its own turn. `message_partition`
 gives one contiguous `[start, end)` span per message that together cover every token, which is what
 mean-pooling activations per turn needs:
 

{interp_engine-1.7.2 → interp_engine-1.7.3}/interp_engine/tokenize.py

@@ -371,6 +371,10 @@ class Tokenize:
         - Within a message block, ``header`` / ``content`` / ``footer`` are separated by rendering
           the same message with EMPTY content and diffing (common prefix = header wrapper, common
           suffix = footer wrapper); the middle is the real content.
+        - A system turn the template injects on its own (Llama's knowledge-cutoff preamble, Qwen2.5's
+          default persona) when the caller sent no system message is tagged ``role="system"`` with
+          ``message_index=None``, so the frontend can show it as its own turn instead of folding it
+          into the first message's header. See :meth:`_injected_system_split`.
         - The trailing ``add_generation_prompt`` scaffold is attributed to the (pending) assistant
           turn as an ``assistant`` ``header`` so the assistant bubble opens before any content.
 
@@ -409,8 +413,8 @@ class Tokenize:
         prefix_end: list[int] = []
         for j in range(0, n + 1):
             if j == 0:
-                # Empty message list: leave any leading scaffold (BOS / system preamble) as part
-                # of message 0's block. Rendering [] would also make some templates raise.
+                # Empty message list: the leading scaffold (BOS / injected system turn) is split
+                # off message 0's block below. Rendering [] would also make some templates raise.
                 prefix_end.append(0)
                 continue
             if rendered is not None:
@@ -438,6 +442,18 @@ class Tokenize:
         msg_idx: list[int | None] = [None] * total
         sections: list[str] = ["scaffold"] * total
 
+        # A system turn the template injected itself sits at the head of message 0's block.
+        # Tag it as its own turn (no message index) and start message 0 after it.
+        injected = None
+        if rendered is None and n > 0:
+            injected = self._injected_system_split(messages, full_ids, prefix_end[1], **template_kwargs)
+        if injected is not None:
+            hdr, cnt, ftr = injected
+            for pos in range(hdr + cnt + ftr):
+                roles[pos] = "system"
+                sections[pos] = "header" if pos < hdr else ("content" if pos < hdr + cnt else "footer")
+            prefix_end[0] = hdr + cnt + ftr
+
         for k in range(n):
             start, end = prefix_end[k], prefix_end[k + 1]
             if end <= start:
@@ -484,6 +500,66 @@ class Tokenize:
             for pos, tid in enumerate(full_ids)
         ]
 
+    def _injected_system_split(
+        self,
+        messages: list[dict[str, str]],
+        full_ids: list[int],
+        block_end: int,
+        **template_kwargs: Any,
+    ) -> tuple[int, int, int] | None:
+        """Return ``(header_len, content_len, footer_len)`` of a system turn the template injected
+        at the head of ``full_ids``, or ``None`` when there is none.
+
+        Only applies when the caller sent no system message. The template is asked to render the
+        same first message behind an explicit system message holding a probe string, and that
+        render is aligned against message 0's real block ``full_ids[:block_end]``:
+
+        - The common prefix is the system header (Llama's knowledge-cutoff preamble lands here,
+          since the template writes it before any system content).
+        - The common suffix is the system footer plus the whole first-message block, which is
+          identical in both renders. The block's own length comes from rendering the system
+          message alone, so the footer is what remains.
+        - Whatever sits between in the real block is the injected content (Qwen2.5's default
+          persona); in the probe render it must be the probe itself, which is what proves the
+          alignment found the system slot and not some other shared token.
+
+        Templates that do not inject a system turn (Qwen3), fold the system message into the
+        first user turn (Gemma-3, Mistral) or reject the system role (Gemma-2) fail one of these
+        checks and get ``None``, leaving message 0's block as it was.
+        """
+        if messages[0].get("role") == "system":
+            return None
+        probe = "system-content-probe"
+        system = {"role": "system", "content": probe}
+        try:
+            sys_only = self._render_ids(
+                [system], add_generation_prompt=False, continue_final_message=False, **template_kwargs
+            )
+            probed = self._render_ids(
+                [system, messages[0]], add_generation_prompt=False, continue_final_message=False, **template_kwargs
+            )
+        except Exception:  # noqa: BLE001 - some templates reject the system role
+            return None
+        if not sys_only or len(sys_only) >= len(probed) or probed[: len(sys_only)] != sys_only:
+            return None
+        block = full_ids[:block_end]
+        first_len = len(probed) - len(sys_only)
+        sys_end = len(block) - first_len
+        if sys_end <= 0:
+            return None
+        suffix = _common_suffix_len(probed, block)
+        if suffix < first_len:
+            return None
+        footer_len = suffix - first_len
+        header_len = min(_common_prefix_len(probed, block), sys_end - footer_len)
+        content_len = sys_end - footer_len - header_len
+        if content_len < 0:
+            return None
+        slot = probed[header_len : len(probed) - suffix]
+        if probe not in self.tokenizer.decode(slot, clean_up_tokenization_spaces=False):
+            return None
+        return header_len, content_len, footer_len
+
     def _header_footer_split(
         self,
         messages: list[dict[str, str]],

{interp_engine-1.7.2 → interp_engine-1.7.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "interp-engine"
-version = "1.7.2"
+version = "1.7.3"
 description = "A fast, standardized, and easy to use interpretability engine."
 readme = "README.md"
 license = "Apache-2.0"

interp_engine-1.7.3/tests/test_injected_system_spans.py

@@ -0,0 +1,115 @@
+"""A system turn the chat template injects on its own gets its own spans.
+
+Llama 3.x writes a knowledge-cutoff preamble and Qwen2.5 a default persona into a system turn
+whenever the caller sends no system message. ``Tokenize.message_spans`` tags that turn
+``role="system"`` with no message index, so a frontend can show it as its own bubble instead of
+folding it into the first user message's header. Templates that inject nothing (Qwen3) or fold
+the system role into the first user turn (Gemma-3) must be left exactly as they were.
+
+Tokenizer-only: nothing here runs a forward, so the (ungated) tokenizer repos are enough.
+"""
+
+from __future__ import annotations
+
+from functools import cache
+
+import pytest
+from transformers import AutoTokenizer
+
+from interp_engine.tokenize import Tokenize, TokenSpan
+
+LLAMA = "unsloth/Llama-3.2-1B-Instruct"
+QWEN25 = "Qwen/Qwen2.5-0.5B-Instruct"
+QWEN3 = "Qwen/Qwen3-0.6B"
+GEMMA3 = "unsloth/gemma-3-1b-it"
+
+USER = [{"role": "user", "content": "What is 2+2?"}]
+SYSTEM_THEN_USER = [{"role": "system", "content": "Be terse."}, *USER]
+
+
+@cache
+def _tok(repo: str) -> Tokenize:
+    try:
+        return Tokenize(AutoTokenizer.from_pretrained(repo))
+    except Exception as exc:  # noqa: BLE001 - offline / uncached
+        pytest.skip(f"could not load tokenizer {repo}: {exc}")
+
+
+def _injected(spans: list[TokenSpan]) -> list[TokenSpan]:
+    return [s for s in spans if s.role == "system" and s.message_index is None]
+
+
+def _text(spans: list[TokenSpan]) -> str:
+    return "".join(s.token_str for s in spans)
+
+
+@pytest.mark.parametrize("repo", [LLAMA, QWEN25])
+def test_injected_system_turn_is_its_own_span_run(repo: str):
+    tok = _tok(repo)
+    spans = tok.message_spans(USER, add_generation_prompt=True)
+    injected = _injected(spans)
+    assert injected, "the template injects a system turn, so one must be tagged"
+    # It is one contiguous run at the very start, closed by a footer.
+    assert [s.position for s in injected] == list(range(len(injected)))
+    assert injected[0].section == "header"
+    assert injected[-1].section == "footer"
+    # Message 0 is the user turn, and it starts right after the injected turn.
+    user = [s for s in spans if s.message_index == 0]
+    assert user and all(s.role == "user" for s in user)
+    assert user[0].position == len(injected)
+    assert "2+2" in _text([s for s in user if s.section == "content"]).replace(" ", "")
+
+
+def test_llama_preamble_is_the_system_header():
+    """Llama writes the knowledge-cutoff lines before any system content, so they are header."""
+    spans = _tok(LLAMA).message_spans(USER, add_generation_prompt=True)
+    injected = _injected(spans)
+    header = _text([s for s in injected if s.section == "header"])
+    assert "Cutting Knowledge Date" in header
+    assert not any(s.section == "content" for s in injected)
+    assert _text([s for s in injected if s.section == "footer"]) == "<|eot_id|>"
+
+
+def test_qwen25_default_persona_is_the_system_content():
+    spans = _tok(QWEN25).message_spans(USER, add_generation_prompt=True)
+    injected = _injected(spans)
+    assert _text([s for s in injected if s.section == "header"]) == "<|im_start|>system\n"
+    assert _text([s for s in injected if s.section == "content"]).startswith("You are Qwen")
+    assert _text([s for s in injected if s.section == "footer"]) == "<|im_end|>\n"
+
+
+@pytest.mark.parametrize("repo", [LLAMA, QWEN25, QWEN3, GEMMA3])
+def test_explicit_system_message_is_message_zero(repo: str):
+    """A caller-supplied system message keeps its index; nothing is tagged as injected."""
+    spans = _tok(repo).message_spans(SYSTEM_THEN_USER, add_generation_prompt=True)
+    assert not _injected(spans)
+    assert any(s.message_index == 0 for s in spans)
+    assert any(s.message_index == 1 and s.role == "user" for s in spans)
+
+
+@pytest.mark.parametrize("repo", [QWEN3, GEMMA3])
+def test_templates_without_an_injected_turn_are_unchanged(repo: str):
+    spans = _tok(repo).message_spans(USER, add_generation_prompt=True)
+    assert not _injected(spans)
+    assert spans[0].message_index == 0
+    assert spans[0].role == "user"
+
+
+@pytest.mark.parametrize("repo", [LLAMA, QWEN25, QWEN3, GEMMA3])
+def test_spans_still_align_with_the_tokenized_prompt(repo: str):
+    tok = _tok(repo)
+    for messages in (USER, SYSTEM_THEN_USER):
+        spans = tok.message_spans(messages, add_generation_prompt=True)
+        ids = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=True)
+        assert [s.token_id for s in spans] == list(ids)
+        assert [s.position for s in spans] == list(range(len(spans)))
+
+
+def test_injected_turn_survives_a_prefill():
+    """The split only touches the head of the sequence; an open final turn stays open."""
+    tok = _tok(LLAMA)
+    messages = [*USER, {"role": "assistant", "content": "It is"}]
+    spans = tok.message_spans(messages, add_generation_prompt=False, continue_final_message=True)
+    assert _injected(spans)
+    prefill = [s for s in spans if s.message_index == 1]
+    assert prefill and prefill[-1].section == "content"