Skip to content

Implement inject_span="trial" in RepengGenerator - #50

Merged
bamdadd merged 1 commit into
bamdadd:mainfrom
mhmsvv:trial-scoped-injection-span
Jul 31, 2026
Merged

Implement inject_span="trial" in RepengGenerator#50
bamdadd merged 1 commit into
bamdadd:mainfrom
mhmsvv:trial-scoped-injection-span

Conversation

@mhmsvv

@mhmsvv mhmsvv commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Implements inject_span="trial" in RepengGenerator: injection now applies only from the newline immediately before "Trial 1" onward (prompt tail + generated response), matching the paper's stricter protocol, instead of only supporting "full" (whole forward pass).
  • Adds resolve_inject_span_start(tokenizer, prompt, inject_span), which computes the anchor token index for "trial" and returns 0 for "full"; unknown span names still raise NotImplementedError.
  • generate_batch now forwards the pre-anchor preamble with control off (to build its KV cache), then continues generation from the anchor with control on — reusing the same cached-generation path as "full" for the anchor-onward span. "full" behavior is unchanged (same single-call code path as before).

Test plan

  • New fast unit tests for resolve_inject_span_start using a fake offset-mapping tokenizer ("full" → token 0; "trial" → lands on the anchor newline; unknown span → NotImplementedError).
  • New fast test asserting RepengGenerator still rejects unknown inject_span values without loading a model.
  • New @pytest.mark.slow test on the real Qwen/Qwen2.5-0.5B-Instruct tokenizer confirming the computed index lands on the newline before "Trial 1" in the rendered prompt, and that "full" still resolves to token 0.
  • uv run pytest -q (full suite, including slow) passes.
  • uv run ruff check / uv run ruff format --check / uv run mypy src/introspection_scaling/harness.py all pass.

Closes #11

🤖 Generated with Claude Code

Injects only from the newline before "Trial 1" onward (prompt tail +
response), matching the paper's stricter span, instead of the whole
forward pass. The preamble is forwarded with control off to build its
KV cache, then generation continues from the anchor with control on.

Adds resolve_inject_span_start() to compute the anchor token index and
tests covering both "full" (token 0) and "trial" (anchor newline), on
a fake offset-mapping tokenizer and on the real Qwen2.5-0.5B-Instruct
tokenizer.

Closes bamdadd#11

Co-Authored-By: Claude Sonnet 5 <[email protected]>
@bamdadd
bamdadd merged commit aeddb4c into bamdadd:main Jul 31, 2026
1 check passed
bamdadd pushed a commit that referenced this pull request Jul 31, 2026
Injects only from the newline before "Trial 1" onward (prompt tail +
response), matching the paper's stricter span, instead of the whole
forward pass. The preamble is forwarded with control off to build its
KV cache, then generation continues from the anchor with control on.

Adds resolve_inject_span_start() to compute the anchor token index and
tests covering both "full" (token 0) and "trial" (anchor newline), on
a fake offset-mapping tokenizer and on the real Qwen2.5-0.5B-Instruct
tokenizer.

Closes #11
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Implement the paper's trial-scoped injection span in RepengGenerator

2 participants