Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,91 @@
+---
+language: en
+license: apache-2.0
+library_name: transformers
+tags:
+- recallm
+- long-context
+- reasoning
+- constrained-decoding
+- information-retrieval
+- in-context-retrieval
+base_model: Qwen/Qwen2.5-7B-Instruct
+datasets:
+- kswhitecross/RecaLLM-data
+pipeline_tag: text-generation
+---
+
+# RecaLLM-Qwen2.5-7B
+
+**RecaLLM** is a family of reasoning language models that interleave reasoning with explicit, constrained-decoding **recall spans** that copy evidence verbatim from context. This addresses the **lost-in-thought** phenomenon: reasoning LLMs lose in-context retrieval ability after chain-of-thought generation.
+
+This model is RecaLLM trained on **Qwen2.5-7B-Instruct**.
+
+**Code:** [https://github.com/kswhitecross/RecaLLM](https://github.com/kswhitecross/RecaLLM)
+
+## Quick Start
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "kswhitecross/RecaLLM-Qwen2.5-7B",
+    trust_remote_code=True,
+    torch_dtype="auto",
+    device_map="auto",
+)
+tokenizer = AutoTokenizer.from_pretrained("kswhitecross/RecaLLM-Qwen2.5-7B")
+
+messages = [{"role": "user", "content": "Your prompt here..."}]
+inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(model.device)
+outputs = model.generate(inputs, max_new_tokens=10240, temperature=0.6, top_p=0.95)
+print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=False))
+```
+
+## Recall Spans
+
+RecaLLM uses special single tokens `<|start_recall|>` and `<|end_recall|>` to delimit recall spans. During a recall span, a constrained decoding mechanism (built into the model class) masks out invalid tokens at each step, allowing only tokens that continue a valid prefix match against the input context. This guarantees that every recall span is a **verbatim contiguous substring** of the searchable context.
+
+The model decides *what* and *when* to recall; constrained decoding ensures the recalled content is exact. Example output:
+
+```
+<think>
+I need to find the population of the city mentioned in the document.
+Looking at the context, I see <|start_recall|>The city of Springfield has a population of 167,376<|end_recall|>.
+So the population is 167,376.
+</think>
+
+Answer: 167,376
+```
+
+## System Prompt
+
+This model was trained with a specific system prompt that instructs it to use recall spans. The chat template automatically includes this prompt when no system message is provided. **Using a different system prompt may significantly degrade performance.** See the [code repository](https://github.com/kswhitecross/RecaLLM) for the full system prompt.
+
+## Inference with vLLM
+
+For vLLM inference, use the `VLLMTokenRecaLLMLogitsProcessor` from the [RecaLLM code repository](https://github.com/kswhitecross/RecaLLM). vLLM should load the model using the **base Qwen2 architecture** (not the custom RecaLLM transformers classes) via `hf_overrides`, with the logits processor applied as an overlay:
+
+```python
+from vllm import LLM, SamplingParams
+from recallm import VLLMTokenRecaLLMLogitsProcessor
+
+llm = LLM(
+    model="kswhitecross/RecaLLM-Qwen2.5-7B",
+    trust_remote_code=True,
+    logits_processors=[VLLMTokenRecaLLMLogitsProcessor],
+    hf_overrides={
+        "model_type": "qwen2",
+        "architectures": ["Qwen2ForCausalLM"],
+    },
+)
+
+sampling_params = SamplingParams(temperature=0.6, top_p=0.95, max_tokens=10240)
+outputs = llm.generate(["Your prompt here..."], sampling_params=sampling_params)
+```
+
+See `recallm/recallm_vllm.py` in the code repo for the full logits processor implementation.
+
+## Citation
+
+<!-- TODO: Add citation once paper is on arXiv -->

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|end_recall|>": 151666,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|start_recall|>": 151665,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a careful reasoning assistant, equipped with a recall tool which lets you\nlook back and faithfully copy information from your context.\n\nFor each task:\n- Restate the goal in your own words.\n- Reason step by step, keeping track of assumptions.\n- If you notice a mistake or contradiction, backtrack and fix it.\n- Verify the final result (e.g., re-check arithmetic, test edge cases,\n  confirm units/constraints).\n- When it would be helpful, use your recall tool with <|start_recall|> ... <|end_recall|> to\n  look back and find relevant information. Be selective — focus on the passages\n  most relevant to your answer.\n\nOutput:\n - First, do all your reasoning and thinking inside of <think> ... </think>\n - After thinking, give a brief justification of your answer, then\n - End with a single final line of the form: Answer: (your final answer here).' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are a careful reasoning assistant, equipped with a recall tool which lets you\nlook back and faithfully copy information from your context.\n\nFor each task:\n- Restate the goal in your own words.\n- Reason step by step, keeping track of assumptions.\n- If you notice a mistake or contradiction, backtrack and fix it.\n- Verify the final result (e.g., re-check arithmetic, test edge cases,\n  confirm units/constraints).\n- When it would be helpful, use your recall tool with <|start_recall|> ... <|end_recall|> to\n  look back and find relevant information. Be selective — focus on the passages\n  most relevant to your answer.\n\nOutput:\n - First, do all your reasoning and thinking inside of <think> ... </think>\n - After thinking, give a brief justification of your answer, then\n - End with a single final line of the form: Answer: (your final answer here).<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n<think>\n' }}
+{%- endif %}

config.json

@@ -0,0 +1,87 @@
+{
+  "architectures": [
+    "RecaLLMQwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "modelling_recallm.RecaLLMQwen2Config",
+    "AutoModel": "modelling_recallm.RecaLLMQwen2Model",
+    "AutoModelForCausalLM": "modelling_recallm.RecaLLMQwen2ForCausalLM"
+  },
+  "close_episode_K_min": 16,
+  "close_episode_eps": 0.2,
+  "close_episode_prob": 0.3,
+  "dtype": "bfloat16",
+  "enable_exp_episodes": false,
+  "eos_token_id": 151645,
+  "force_recall": true,
+  "hidden_act": "silu",
+  "hidden_size": 3584,
+  "initializer_range": 0.02,
+  "intermediate_size": 18944,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_recall_length": 128,
+  "max_window_layers": 28,
+  "model_type": "recallm_qwen2",
+  "no_episode_prob": 0.4,
+  "num_attention_heads": 28,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 4,
+  "open_episode_K_max": 32,
+  "open_episode_eps": 0.6,
+  "open_episode_prob": 0.3,
+  "pad_token_id": 151643,
+  "profile_recall_masking": false,
+  "recall_end_str": "</recall>",
+  "recall_end_token_id": 151666,
+  "recall_start_str": "<recall>",
+  "recall_start_token_id": 151665,
+  "recall_tag_type": "token_id",
+  "record_recall_prob_after_generate": false,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": {
+    "factor": 4.0,
+    "original_max_position_embeddings": 32768,
+    "rope_type": "yarn",
+    "type": "yarn"
+  },
+  "rope_theta": 1000000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "use_fast_recall": true,
+  "use_sliding_window": false,
+  "vllm_disable_recall_masking": false,
+  "vocab_size": 151669
+}

generation_config.json

@@ -0,0 +1,6 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "transformers_version": "4.57.6"
+}

model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9f4af3bca36404999e7084f9219c76c82f6d8a239f112d9d19d17081bed5c96
+size 4999665008

model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4fa11859ea998f992e77230169631f6a768a4bd97ba3bbdae27bd0ffa06e71d
+size 4866690688

model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe95fc040ef0a97743e0d28b020988703bb6dc11540b2797c4df9f786c4aed99
+size 4984886976

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a725a7fe1a91a27d3740a54bb8da579ff9da9b2e1c28ef1a2363d35976363fa
+size 374366472

model.safetensors.index.json

@@ -0,0 +1,347 @@
+{
+  "metadata": {
+    "total_parameters": 7612785152,
+    "total_size": 15225570304
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00001-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00003-of-00004.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.bias": "model-00004-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.bias": "model-00004-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.bias": "model-00004-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.bias": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.bias": "model-00003-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.bias": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.norm.weight": "model-00002-of-00004.safetensors"
+  }
+}

modelling_recallm.py

@@ -0,0 +1,604 @@
+"""
+RecaLLM model classes and recall masking logic.
+
+Provides the RecaLLM mixin that adds constrained-decoding recall spans to
+any causal language model. During a recall span, the model can only generate
+tokens that continue a valid prefix match against the input context.
+
+Supports Llama and Qwen2 architectures.
+"""
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    AutoModel,
+    AutoTokenizer,
+    LlamaConfig,
+    LlamaForCausalLM,
+    LlamaModel,
+    Qwen2Config,
+    Qwen2ForCausalLM,
+    Qwen2Model,
+)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+import torch
+from typing import Optional, Union
+import contextlib
+import inspect
+import threading
+
+# Large finite negative used instead of float('-inf') for recall masking.
+# Using -inf triggers VeRL's isinf() guard in logprobs_from_logits(), which
+# diverts from the fast Flash Attention Triton kernel to a row-by-row Python
+# loop (~60x slower). exp(-1e4) underflows to exactly 0 in bfloat16/float16/float32,
+# so softmax behavior is identical, but isinf() returns False.
+# -1e4 is the largest safe value across all dtypes (float16 max is 65504;
+# -1e5 overflows to -inf in float16).
+RECALL_MASK_VALUE = -1e4
+
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class RecallSpan:
+    """Tracks a single recall span within a sequence during generation."""
+    input_idx_start: int
+    input_idx_end: int
+    logits_start: int
+    logits_end: int
+    span_start: int
+    span_end: int
+    allowed_tokens_list: list = field(default_factory=list)
+
+    def get_mask(self, vocab_size: int) -> torch.Tensor:
+        """Build a [span_length, vocab_size] bool mask (True = disallowed)."""
+        n_positions = len(self.allowed_tokens_list)
+        mask = torch.ones(n_positions, vocab_size, dtype=torch.bool)
+        for i, allowed in enumerate(self.allowed_tokens_list):
+            mask[i, allowed] = False
+        return mask
+
+
+def get_prefix_continuations(
+    context: torch.Tensor,
+    prefix: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Find all valid next-token continuations of ``prefix`` in ``context``.
+
+    Parameters
+    ----------
+    context : torch.Tensor
+        1-D tensor of token IDs representing the available context.
+    prefix : torch.Tensor
+        1-D tensor of token IDs representing the current recall span so far.
+
+    Returns
+    -------
+    torch.Tensor
+        1-D tensor of unique token IDs that could validly continue the prefix.
+    """
+    n = len(prefix)
+    if n == 0:
+        return context.unique()
+    if n >= len(context):
+        return torch.tensor([], dtype=context.dtype, device=context.device)
+
+    windows = context.unfold(0, n, 1)
+    match_mask = (windows == prefix.unsqueeze(0)).all(dim=1)
+    match_indices = match_mask.nonzero(as_tuple=False).squeeze(1)
+    continuation_indices = match_indices + n
+    valid = continuation_indices < len(context)
+    continuation_indices = continuation_indices[valid]
+
+    if len(continuation_indices) == 0:
+        return torch.tensor([], dtype=context.dtype, device=context.device)
+    return context[continuation_indices].unique()
+
+
+def _get_prefix_continuations_incremental(
+    context: torch.Tensor,
+    span_tokens: torch.Tensor,
+    recall_start_token_id: int,
+    recall_end_token_id: int,
+) -> list[torch.Tensor]:
+    """
+    Compute valid continuations for every prefix length in a recall span.
+
+    Returns R+1 tensors (for the empty prefix through the full span), each
+    containing the allowed token IDs at that position. The recall start token
+    is excluded from continuations and the recall end token is always appended.
+
+    This is O(C) per span (single pass over context) instead of O(C*R).
+    """
+    C = len(context)
+    R = len(span_tokens)
+    device = context.device
+
+    end_tok = torch.tensor([recall_end_token_id], device=device)
+
+    if C == 0:
+        return [end_tok] * (R + 1)
+
+    # Build match matrix: match[r][c] = True iff span_tokens[r] == context[c]
+    if R > 0:
+        match = (span_tokens.unsqueeze(1) == context.unsqueeze(0))
+    else:
+        match = torch.empty(0, C, dtype=torch.bool, device=device)
+
+    # alive[c] tracks which context positions still match the prefix at each step
+    alive = torch.ones(C, dtype=torch.bool, device=device)
+    results = []
+
+    # k=0: empty prefix -> any context token is a valid continuation
+    conts_0 = context[alive].unique()
+    conts_0 = conts_0[conts_0 != recall_start_token_id]
+    results.append(torch.cat([conts_0, end_tok]))
+
+    for k in range(R):
+        alive_shifted = torch.zeros(C, dtype=torch.bool, device=device)
+        if k == 0:
+            alive_shifted[1:] = match[0, :-1]
+        else:
+            prev_alive = alive
+            alive_shifted[1:] = prev_alive[:-1] & match[k, :-1]
+        alive = alive_shifted
+
+        if not alive.any():
+            results.append(end_tok)
+        else:
+            conts = context[alive].unique()
+            conts = conts[conts != recall_start_token_id]
+            results.append(torch.cat([conts, end_tok]))
+
+    return results
+
+
+class RecaLLMConfigMixin:
+    """
+    Configuration mixin for RecaLLM models.
+
+    Parameters
+    ----------
+    force_recall : bool, default True
+        If True, mask out logits for tokens that would produce non-viable
+        sequences inside recall spans.
+    recall_tag_type : str, default "token_id"
+        Must be "token_id" — uses single-token delimiters for recall spans.
+    recall_start_token_id : int
+        Token ID for the start of a recall span.
+    recall_end_token_id : int
+        Token ID for the end of a recall span.
+    recall_start_str : str, default "<recall>"
+        String representation of recall start (for display/prompts).
+    recall_end_str : str, default "</recall>"
+        String representation of recall end (for display/prompts).
+    max_recall_length : int, default 128
+        Maximum tokens in a recall span before forcing closure.
+    use_fast_recall : bool, default True
+        Use the optimized background-thread forward pass.
+    """
+    def __init__(
+            self,
+            force_recall: bool = True,
+            recall_tag_type: str = "token_id",
+            recall_start_str: Optional[str] = "<recall>",
+            recall_end_str: Optional[str] = "</recall>",
+            recall_start_token_id: Optional[int] = None,
+            recall_end_token_id: Optional[int] = None,
+            record_recall_prob_after_generate: bool = False,
+            max_recall_length: int = 128,
+            use_fast_recall: bool = True,
+            **kwargs):
+        super().__init__(**kwargs)
+        self.force_recall = force_recall
+        self.recall_tag_type = recall_tag_type
+        self.recall_start_str = recall_start_str
+        self.recall_end_str = recall_end_str
+        self.recall_start_token_id = recall_start_token_id
+        self.recall_end_token_id = recall_end_token_id
+        self.record_recall_prob_after_generate = record_recall_prob_after_generate
+        self.max_recall_length = max_recall_length
+        self.use_fast_recall = use_fast_recall
+
+
+class RecaLLMMixin:
+    """
+    Mixin that adds RecaLLM recall masking to any causal LM.
+
+    Overrides ``forward()`` to mask logits at positions inside recall spans,
+    allowing only tokens that continue a valid prefix match against the
+    input context. Uses background threading + CUDA stream overlap for
+    efficient masking during training.
+    """
+    def __init__(self, config: RecaLLMConfigMixin, recall_enabled_in_forward_pass: bool = True):
+        super().__init__(config)
+        self.recall_enabled_in_forward_pass = recall_enabled_in_forward_pass
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            logits_to_keep: Union[int, torch.LongTensor] = 0,
+            recall_enabled: bool | None = None,
+            **kwargs
+        ) -> CausalLMOutputWithPast:
+        """
+        Forward pass with optional recall masking.
+
+        When recall masking is enabled, logits at positions inside recall spans
+        are masked so that only valid context continuations are allowed.
+
+        Parameters
+        ----------
+        recall_enabled : bool or None
+            Override for recall masking. None uses ``self.recall_enabled_in_forward_pass``.
+        """
+        if recall_enabled is None:
+            recall_enabled = self.recall_enabled_in_forward_pass
+
+        if recall_enabled:
+            assert input_ids is not None, \
+                "`input_ids` must be provided when recall masking is enabled."
+            assert self.config.recall_tag_type == "token_id", \
+                "Only recall_tag_type='token_id' is supported."
+            return self._fast_recall_forward(
+                input_ids=input_ids, attention_mask=attention_mask,
+                position_ids=position_ids, labels=labels,
+                logits_to_keep=logits_to_keep, **kwargs)
+
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            labels=labels,
+            logits_to_keep=logits_to_keep,
+            **kwargs,
+        )
+
+    @staticmethod
+    def _compute_recall_masks(
+        input_ids_cpu: torch.Tensor,
+        attention_mask_cpu: Optional[torch.Tensor],
+        position_ids_cpu: Optional[torch.Tensor],
+        sequence_length: int,
+        vocab_size: int,
+        start_id: int,
+        end_id: int,
+        eos_ids_list: list[int],
+        keep_logit_start: Optional[int],
+        logits_to_keep_map: Optional[dict],
+    ) -> list[tuple[int, int, int, torch.Tensor]]:
+        """
+        Compute all recall masks on CPU (runs in a background thread).
+
+        Detects recall spans via a state machine, computes valid prefix
+        continuations for each position, and builds boolean masks.
+
+        Returns
+        -------
+        list of (batch_idx, logits_start, logits_end, bool_mask_cpu) tuples.
+        """
+        batch_size = input_ids_cpu.shape[0]
+        mask_instructions = []
+
+        for batch_idx in range(batch_size):
+            seq_cpu = input_ids_cpu[batch_idx]
+            attn_cpu = attention_mask_cpu[batch_idx] if attention_mask_cpu is not None else None
+            pos_cpu = position_ids_cpu[batch_idx] if position_ids_cpu is not None else None
+
+            # Determine sequence bounds
+            if attn_cpu is not None:
+                nonzero = attn_cpu.nonzero(as_tuple=False).squeeze(1)
+                start_idx = nonzero.min().item()
+                end_idx = nonzero.max().item() + 1
+            else:
+                start_idx = 0
+                end_idx = sequence_length
+
+            # Vectorized event detection
+            seq_slice = seq_cpu[start_idx:end_idx]
+
+            is_start = (seq_slice == start_id)
+            start_positions = is_start.nonzero(as_tuple=False).squeeze(1) + start_idx
+
+            is_end = (seq_slice == end_id)
+            for eid in eos_ids_list:
+                is_end = is_end | (seq_slice == eid)
+            end_positions = is_end.nonzero(as_tuple=False).squeeze(1) + start_idx
+
+            if pos_cpu is not None:
+                pos_slice = pos_cpu[start_idx:end_idx]
+                is_boundary = torch.zeros(pos_slice.shape[0], dtype=torch.bool)
+                if pos_slice.shape[0] > 1:
+                    is_boundary[1:] = (pos_slice[1:] == 0)
+                boundary_positions = is_boundary.nonzero(as_tuple=False).squeeze(1) + start_idx
+            else:
+                boundary_positions = torch.tensor([], dtype=torch.long)
+
+            # Build sorted event list
+            events = []
+            for p in start_positions.tolist():
+                events.append((p, 0))  # start recall
+            for p in end_positions.tolist():
+                events.append((p, 1))  # end recall
+            for p in boundary_positions.tolist():
+                events.append((p, 2))  # sequence boundary
+            events.sort(key=lambda e: (e[0], e[1]))
+
+            # State machine to extract recall spans
+            in_recall_mode = False
+            context_start = start_idx
+            current_content_start = None
+            current_ctx_start = None
+            raw_spans = []
+
+            for pos, etype in events:
+                if etype == 0:
+                    if not in_recall_mode:
+                        in_recall_mode = True
+                        current_content_start = pos + 1
+                        current_ctx_start = context_start
+                elif etype == 1:
+                    if in_recall_mode:
+                        in_recall_mode = False
+                        raw_spans.append((current_ctx_start, current_content_start, pos))
+                        current_content_start = None
+                        current_ctx_start = None
+
+                if etype == 2:
+                    if in_recall_mode:
+                        raw_spans.append((current_ctx_start, current_content_start, pos))
+                        in_recall_mode = False
+                        current_content_start = None
+                        current_ctx_start = None
+                    context_start = pos
+
+            if in_recall_mode:
+                raw_spans.append((current_ctx_start, current_content_start, end_idx))
+
+            # For each span: prefix match + build bool mask
+            for (ctx_start, content_start, content_end) in raw_spans:
+                if content_end < content_start:
+                    continue
+
+                start_token_pos = content_start - 1
+
+                if keep_logit_start is not None:
+                    if content_end - 1 < keep_logit_start:
+                        continue
+                elif logits_to_keep_map is not None:
+                    has_any = False
+                    for p in range(start_token_pos, content_end):
+                        if p in logits_to_keep_map:
+                            has_any = True
+                            break
+                    if not has_any:
+                        continue
+
+                context = seq_cpu[ctx_start:content_start - 1]
+                span_tokens = seq_cpu[content_start:content_end]
+
+                allowed_lists = _get_prefix_continuations_incremental(
+                    context=context,
+                    span_tokens=span_tokens,
+                    recall_start_token_id=start_id,
+                    recall_end_token_id=end_id,
+                )
+
+                # Filter to positions that have logits
+                mask_positions = list(range(start_token_pos, content_end))
+                filtered_allowed = []
+                logits_start = None
+                logits_end = None
+
+                for i, abs_pos in enumerate(mask_positions):
+                    logits_idx = None
+                    if keep_logit_start is not None:
+                        if abs_pos >= keep_logit_start:
+                            logits_idx = abs_pos - keep_logit_start
+                    elif logits_to_keep_map is not None:
+                        if abs_pos in logits_to_keep_map:
+                            logits_idx = logits_to_keep_map[abs_pos]
+
+                    if logits_idx is not None:
+                        filtered_allowed.append(allowed_lists[i])
+                        if logits_start is None:
+                            logits_start = logits_idx
+                        logits_end = logits_idx + 1
+
+                # Build boolean mask on CPU
+                if logits_start is not None:
+                    N = len(filtered_allowed)
+                    bool_mask = torch.ones(N, vocab_size, dtype=torch.bool, pin_memory=True)
+                    sizes = torch.tensor([a.numel() for a in filtered_allowed])
+                    row_idx = torch.repeat_interleave(torch.arange(N), sizes)
+                    col_idx = torch.cat(filtered_allowed)
+                    bool_mask[row_idx, col_idx] = False
+
+                    mask_instructions.append((batch_idx, logits_start, logits_end, bool_mask))
+
+        return mask_instructions
+
+    def _fast_recall_forward(
+            self,
+            input_ids: torch.LongTensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            logits_to_keep: Union[int, torch.LongTensor] = 0,
+            **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """
+        Optimized forward pass with recall masking (token_id mode only).
+
+        Overlaps CPU-bound mask computation and GPU mask transfer with the
+        model forward pass using a background thread + CUDA stream.
+
+        Flow:
+          1. Copy input_ids / attention_mask / position_ids to CPU
+          2. Background thread: span detection + prefix matching + bool mask
+             building (CPU), then transfer masks to GPU on a dedicated stream
+          3. Run super().forward() on the default CUDA stream (overlaps with 2)
+          4. Join thread, sync streams, apply masks to logits
+          5. Compute loss if labels provided
+        """
+        assert self.config.recall_tag_type == "token_id", \
+            "_fast_recall_forward only supports recall_tag_type='token_id'"
+
+        batch_size = input_ids.shape[0]
+        sequence_length = input_ids.shape[1]
+
+        # --- Step A: Precompute config values ---
+        keep_logit_start = None
+        if isinstance(logits_to_keep, int):
+            keep_logit_start = 0 if logits_to_keep == 0 else sequence_length - logits_to_keep
+
+        logits_to_keep_map = None
+        if isinstance(logits_to_keep, torch.Tensor):
+            ltk_cpu = logits_to_keep.cpu()
+            logits_to_keep_map = {pos.item(): idx for idx, pos in enumerate(ltk_cpu)}
+
+        start_id = self.config.recall_start_token_id
+        end_id = self.config.recall_end_token_id
+        eos_id = self.config.eos_token_id
+        if isinstance(eos_id, int):
+            eos_ids_list = [eos_id]
+        elif isinstance(eos_id, list):
+            eos_ids_list = eos_id
+        else:
+            eos_ids_list = []
+
+        vocab_size = self.config.vocab_size
+
+        # --- Step B: Copy inputs to CPU ---
+        input_ids_cpu = input_ids.cpu()
+        attention_mask_cpu = attention_mask.cpu() if attention_mask is not None else None
+        position_ids_cpu = position_ids.cpu() if position_ids is not None else None
+
+        # --- Step C: Background thread for CPU mask computation + GPU transfer ---
+        target_device = input_ids.device
+        thread_result = [None]
+        thread_stream = [None]
+        thread_error = [None]
+
+        def _mask_thread():
+            try:
+                masks_cpu = self._compute_recall_masks(
+                    input_ids_cpu=input_ids_cpu,
+                    attention_mask_cpu=attention_mask_cpu,
+                    position_ids_cpu=position_ids_cpu,
+                    sequence_length=sequence_length,
+                    vocab_size=vocab_size,
+                    start_id=start_id,
+                    end_id=end_id,
+                    eos_ids_list=eos_ids_list,
+                    keep_logit_start=keep_logit_start,
+                    logits_to_keep_map=logits_to_keep_map,
+                )
+                if masks_cpu and target_device.type == 'cuda':
+                    stream = torch.cuda.Stream(device=target_device)
+                    masks_gpu = []
+                    with torch.cuda.stream(stream):
+                        for batch_idx, ls, le, bm_cpu in masks_cpu:
+                            bm_gpu = bm_cpu.to(target_device, non_blocking=True)
+                            masks_gpu.append((batch_idx, ls, le, bm_gpu))
+                    thread_result[0] = masks_gpu
+                    thread_stream[0] = stream
+                else:
+                    thread_result[0] = masks_cpu
+            except Exception as e:
+                thread_error[0] = e
+
+        mask_thread = threading.Thread(target=_mask_thread, daemon=True)
+        mask_thread.start()
+
+        # --- Step D: Run base model forward on GPU (overlaps with thread) ---
+        outputs = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            logits_to_keep=logits_to_keep,
+            position_ids=position_ids,
+            **kwargs,
+        )
+
+        # --- Step E: Wait for background thread ---
+        mask_thread.join()
+
+        if thread_error[0] is not None:
+            raise thread_error[0]
+
+        mask_instructions = thread_result[0]
+        logits_device = outputs.logits.device
+
+        # --- Step F: Apply masks to logits ---
+        if thread_stream[0] is not None:
+            torch.cuda.current_stream(logits_device).wait_stream(thread_stream[0])
+
+        for batch_idx, logits_start, logits_end, bool_mask in mask_instructions:
+            outputs.logits[batch_idx, logits_start:logits_end, :].masked_fill_(
+                bool_mask, RECALL_MASK_VALUE
+            )
+
+        # --- Step G: Compute loss if labels provided ---
+        if labels is not None:
+            loss_func_params = set(inspect.signature(self.loss_function).parameters.keys())
+            loss_func_kwargs = {key: val for key, val in kwargs.items() if key in loss_func_params}
+            outputs["loss"] = self.loss_function(
+                logits=outputs.logits, labels=labels,
+                vocab_size=self.config.vocab_size, **loss_func_kwargs,
+            )
+
+        return outputs
+
+    def enable_recall_in_forward_pass(self):
+        self.recall_enabled_in_forward_pass = True
+
+    def disable_recall_in_forward_pass(self):
+        self.recall_enabled_in_forward_pass = False
+
+    @contextlib.contextmanager
+    def temp_disable_recall_in_forward_pass(self):
+        original = self.recall_enabled_in_forward_pass
+        self.disable_recall_in_forward_pass()
+        try:
+            yield
+        finally:
+            self.recall_enabled_in_forward_pass = original
+
+    def get_base_model(self):
+        return super()
+
+
+# ── Model Registration ──────────────────────────────────────────────────────
+# Llama 3 RecaLLM models
+class RecaLLMLlamaConfig(RecaLLMConfigMixin, LlamaConfig):
+    model_type = "recallm_llama"
+
+class RecaLLMLlamaForCausalLM(RecaLLMMixin, LlamaForCausalLM):
+    config_class = RecaLLMLlamaConfig
+
+class RecaLLMLlamaModel(RecaLLMMixin, LlamaModel):
+    config_class = RecaLLMLlamaConfig
+    _supports_attention_backend = True
+
+AutoConfig.register("recallm_llama", RecaLLMLlamaConfig)
+AutoModelForCausalLM.register(RecaLLMLlamaConfig, RecaLLMLlamaForCausalLM)
+AutoModel.register(RecaLLMLlamaConfig, RecaLLMLlamaModel)
+
+# Qwen 2 (and 2.5) RecaLLM models
+class RecaLLMQwen2Config(RecaLLMConfigMixin, Qwen2Config):
+    model_type = "recallm_qwen2"
+
+class RecaLLMQwen2ForCausalLM(RecaLLMMixin, Qwen2ForCausalLM):
+    config_class = RecaLLMQwen2Config
+
+class RecaLLMQwen2Model(RecaLLMMixin, Qwen2Model):
+    config_class = RecaLLMQwen2Config
+    _supports_attention_backend = True
+
+AutoConfig.register("recallm_qwen2", RecaLLMQwen2Config)
+AutoModelForCausalLM.register(RecaLLMQwen2Config, RecaLLMQwen2ForCausalLM)
+AutoModel.register(RecaLLMQwen2Config, RecaLLMQwen2Model)

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1bd65c75b9f0aad032a905964e080ce1964b14fe101883d40fe86c5a39624f6c
+size 11422653

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<|start_recall|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "<|end_recall|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}