# ThinkPack

> A framework for training, parsing, and evaluating explicit reasoning models — focussed on reasoning-trace collapse.

`thinkpack` provides four core modules for measuring and mitigating reasoning-trace collapse in fine-tuned language models, plus distillation helpers.
Install from PyPI:

```
pip install thinkpack
```

## What is reasoning-trace collapse?

When reasoning models are fine-tuned on standard instruction-response data (answers with no reasoning traces), they can stop producing valid reasoning traces — this is *reasoning-trace collapse*. The model may still answer correctly, but its reasoning block ends up empty or missing.

```
before fine-tuning:   x -> <think>reasoning</think> answer
after naive SFT:      x -> <think></think> answer
or simply:            x -> answer
```

ThinkPack makes this measurable, and provides loss masking as a mitigation (its effect is model- and task-dependent, so measure VR after training).

## What are reasoning blocks?

Reasoning models produce outputs with a reasoning block followed by the final answer:

```
<think>step-by-step reasoning...</think>
final answer
```

Different models format this differently, and `thinkpack` handles the differences automatically.

## Model detection

Every function detects the model's reasoning format from the tokenizer's chat template, and caches the result. `detect_model(tokenizer)` returns a `ModelInfo`:

- `tag_content` and `tag_style` — the reasoning tag. `TagStyle` (a `StrEnum`) is `HTML` for `<think>...</think>` (most models) or `BRACKET` for `[THINK]...[/THINK]` (e.g. Mistral). `open_tag` and `close_tag` give the full tag strings.
- `prefixed` — true if the template adds the opening reasoning tag to the generation prompt (e.g. OLMo-3, DeepSeek-R1).
- `strips_think_tags` — true if the template strips reasoning from the final assistant message (e.g. DeepSeek-R1).
- `strips_history_think_tags` — true if the template strips reasoning from earlier assistant messages (e.g. Qwen3, Qwen3.5).

Pass `override_tag=` (e.g. `"<reasoning>"` or `"[THINK]"`) to any function if the wrong tag is detected, or pass a custom `ModelInfo` as `model_info=` to skip detection entirely (`parse()` then needs no tokenizer). `prefixed` is only used by `parse()`, as templating checks the rendered prompt directly. Multimodal processors (e.g. Qwen3.5 via `AutoProcessor`) can be passed wherever a tokenizer is expected.

## Module overview

| Module | Purpose |
|---|---|
| `thinkpack.chat` | Apply chat templates with optional thought-steering and reasoning history embedding — works the same across all models |
| `thinkpack.parse` | Split raw model output into `reasoning` and `answer`, with flags for presence, validity, and truncation |
| `thinkpack.stats` | Aggregate a batch of `ParsedResponse` objects into VR, ER, TR, MR, pass@1, and Rpass@1 — the collapse metrics |
| `thinkpack.mask` | Tokenize training conversations into a HuggingFace dataset with configurable loss masking — can help mitigate collapse during fine-tuning |
| `thinkpack.distill` | Build prompts asking a teacher model for reasoning traces, and add the extracted traces to training records |

## Key metrics

- **VR (Valid Reasoning)** — fraction of responses with complete, non-blank reasoning (`valid_reasoning_rate`, `.vr`)
- **ER (Empty Reasoning)** — reasoning block opened and closed, but blank (`empty_reasoning_rate`, `.er`)
- **TR (Truncated Reasoning)** — reasoning block opened but never closed (`truncated_reasoning_rate`, `.tr`)
- **MR (Missing Reasoning)** — no reasoning block at all (`missing_reasoning_rate`, `.mr`)
- **pass@1** — standard answer accuracy (`pass_at_1`)
- **Rpass@1** — accuracy among responses with valid reasoning (`rpass_at_1`)

VR + ER + TR + MR = 1. Note `.ar` is the answer rate (`answer_rate`), not a reasoning metric. For nested `[task][sample]` input, rates are averaged across tasks; Rpass@1 averages only over tasks with at least one valid-reasoning sample.

Reasoning-trace collapse shows up as VR → 0 over training steps or data size.

## Public API

```python
import thinkpack

# model detection
thinkpack.detect_model(tokenizer)       # -> ModelInfo (cached per chat template)
thinkpack.get_model_info(tokenizer, override_tag=None)  # -> ModelInfo, with an optional tag override
thinkpack.ModelInfo                     # dataclass: prefixed, tag_content, tag_style, strips_think_tags, strips_history_think_tags; properties open_tag, close_tag
thinkpack.TagStyle                      # StrEnum: HTML | BRACKET

# chat templates
thinkpack.apply_chat_template(conversation, tokenizer, add_generation_prompt=True)   # -> str
thinkpack.apply_chat_templates(conversations, tokenizer, add_generation_prompt=True) # -> list[str]
# add_generation_prompt is passed to the tokenizer, which defaults to False — pass True for generation prompts
# assistant messages may include a "reasoning" key, embedded as the think block
# options: think_prefix, response_prefix, override_tag, model_info (skips detection),
#   add_generation_reasoning (None = unchanged | True = ensure open tag | False = remove open tag),
#   add_history_reasoning (None = template decides | True = always keep | False = always drop)
# any other kwargs are passed to tokenizer.apply_chat_template()

# parsing
thinkpack.parse(response, tokenizer)    # -> ParsedResponse | list[ParsedResponse] | list[list[ParsedResponse]], matching the input shape
thinkpack.parse(response, tokenizer, prompt=prompts)  # pass the generation prompts, so parse knows if the output starts inside an open reasoning block
# options: model_info (skips detection), override_tag, calculate_tokens, add_generation_reasoning
thinkpack.ParsedResponse                # dataclass: answer, reasoning, reasoning_tag, has_valid_reasoning, has_empty_reasoning, has_truncated_reasoning, has_missing_reasoning, reasoning_token_count, answer_token_count; properties has_invalid_reasoning, has_answer

# statistics (VR / ER / TR / MR / pass@1 / Rpass@1)
thinkpack.compute_stats(responses, results=None)  # -> ResponseStats; results is optional bools, same shape as responses
thinkpack.ResponseStats                 # dataclass: total, valid_reasoning_rate, invalid_reasoning_rate, missing_reasoning_rate, truncated_reasoning_rate, empty_reasoning_rate, answer_rate, avg_reasoning_tokens, avg_answer_tokens, pass_at_1, rpass_at_1

# training (loss masking — can help mitigate collapse)
thinkpack.apply_mask(conversations, tokenizer, masked=thinkpack.MaskType.THINK)  # -> Dataset (input_ids, labels, attention_mask)
thinkpack.MaskType                      # IntFlag: PROMPT | THINK | RESPONSE, combine with |; masked=None trains on all tokens
# options: max_seq_length=32768, ignore_index=-100, override_tag, model_info (skips detection), add_history_reasoning
# each conversation must end with the assistant message to train on; requires a fast tokenizer
# sequences are not padded — use a collator that pads labels with -100, e.g. DataCollatorForSeq2Seq

# distillation
thinkpack.build_prompts(records)                         # -> list[str], asking for reasoning in <reasoning_steps> tags
thinkpack.extract_distilled_reasoning(text)              # -> str | None | list[str | None]
thinkpack.update_records(records, responses)             # -> list[dict], with "reasoning" added where found
thinkpack.to_conversations(records)                      # -> list[list[dict]], ready for apply_mask or apply_chat_template
```

## Docs

- [chat.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/src/thinkpack/chat.py): chat templates — `apply_chat_template()`, `apply_chat_templates()`
- [parse.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/src/thinkpack/parse.py): response parsing — `ParsedResponse`, `parse()`
- [stats.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/src/thinkpack/stats.py): statistics — `ResponseStats`, `compute_stats()`
- [mask.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/src/thinkpack/mask.py): loss masking — `MaskType`, `apply_mask()`
- [model.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/src/thinkpack/model.py): model detection — `TagStyle`, `ModelInfo`, `detect_model()`, `get_model_info()`
- [distill.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/src/thinkpack/distill.py): distillation — `build_prompts()`, `extract_distilled_reasoning()`, `update_records()`, `to_conversations()`

## Examples

- [examples/scripts/training.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/examples/scripts/training.py): naive SFT vs masking-based SFT
- [examples/scripts/inference.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/examples/scripts/inference.py): measuring reasoning-trace collapse with parse + compute_stats

## Optional

- [tests/test_mask.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/tests/test_mask.py): mask tests across HTML and BRACKET tag styles, and prefixed vs non-prefixed templates
- [tests/test_parse.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/tests/test_parse.py): parsing tests covering all four response formats
- [tests/test_stats.py](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/tests/test_stats.py): stats aggregation tests
- [pyproject.toml](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/pyproject.toml): package metadata and dependencies
- [README.md](https://raw.githubusercontent.com/itsluketwist/thinkpack/main/README.md): narrative documentation with per-module code examples
