add qwen3-4b and mistral-7b (niah working)

chat_templates/Qwen/Qwen3-4B-Instruct-2507.jinja

@@ -0,0 +1,18 @@
+{%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role + '\n' }}{% generation %}{{ content + '<|im_end|>' }}{% endgeneration %}{{ '\n' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

chat_templates/mistralai/Mistral-7B-Instruct-v0.2.jinja

@@ -0,0 +1,24 @@
+{%- if messages[0]['role'] == 'system' %}
+    {%- set system_message = messages[0]['content'] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+
+{{- bos_token }}
+{%- for message in loop_messages %}
+    {%- if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}
+        {{- raise_exception('After the optional system message, conversation roles must alternate user/assistant/user/assistant/...') }}
+    {%- endif %}
+    {%- if message['role'] == 'user' %}
+        {%- if loop.first and system_message is defined %}
+            {{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST] ' }}
+        {%- else %}
+            {{- ' [INST] ' + message['content'] + ' [/INST] ' }}
+        {%- endif %}
+    {%- elif message['role'] == 'assistant' %}
+        {% generation %}{{ message['content'] + eos_token }}{% endgeneration %}
+    {%- else %}
+        {{- raise_exception('Only user and assistant roles are supported, with the exception of an initial optional system message!') }}
+    {%- endif %}
+{%- endfor %}

scripts/niah/train_mistral_7b.sh

@@ -0,0 +1,39 @@
+WANDB_MODE=disabled run uv run train.py configs/niah_exp/ctx_magic_number_32_256.yaml \
+--model_name_or_path=mistralai/Mistral-7B-Instruct-v0.2 \
+--num_train_epochs=1 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=64 \
+--per_device_eval_batch_size=16 \
+--exp_setup=hyper_lora \
+--aggregator_type=perceiver \
+--target_modules=down_proj \
+--num_blocks=8 \
+--num_self_attn_per_block=0 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_steps=100 \
+--logging_steps=10 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=0 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=2.0 \
+--use_sequence_packing=True \
+--max_packed_inp_len=1024 \
+--max_packed_ctx_len=1024 \
+--dataloader_num_workers=0 \
+--dataloader_prefetch_factor=None \
+--eval_on_start=False \
+--ctx_encoder_type=early_exit \
+--n_latent_queries=208 \
+--use_kl_loss=False \
+--eval_on_start=True \
+--max_ctx_chunk_len=512 \
+--min_ctx_chunk_len=25 \
+--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
+--max_val_samples_per_ds=100 \
+--seed=1 \
+--use_per_ctx_average_loss=True \
+--torch_empty_cache_steps=10

scripts/niah/train_qwen3_4b.sh

@@ -0,0 +1,39 @@
+WANDB_MODE=disabled run uv run train.py configs/niah_exp/ctx_magic_number_32_256.yaml \
+--model_name_or_path=Qwen/Qwen3-4B-Instruct-2507 \
+--num_train_epochs=1 \
+--per_device_train_batch_size=-1 \
+--gradient_accumulation_steps=32 \
+--per_device_eval_batch_size=16 \
+--exp_setup=hyper_lora \
+--aggregator_type=perceiver \
+--target_modules=down_proj \
+--num_blocks=8 \
+--num_self_attn_per_block=0 \
+--num_pre_head_layers=1 \
+--lora_r=8 \
+--eval_steps=100 \
+--logging_steps=10 \
+--save_steps=1000 \
+--learning_rate=4e-5 \
+--lora_dropout=0.0 \
+--neftune_noise_alpha=0 \
+--per_rank_gen=True \
+--per_layer_processing=True \
+--gen_lora_l1_reg_coef=0.5 \
+--use_sequence_packing=True \
+--max_packed_inp_len=2048 \
+--max_packed_ctx_len=2048 \
+--dataloader_num_workers=0 \
+--dataloader_prefetch_factor=None \
+--eval_on_start=False \
+--ctx_encoder_type=early_exit \
+--n_latent_queries=208 \
+--use_kl_loss=False \
+--eval_on_start=True \
+--max_ctx_chunk_len=512 \
+--min_ctx_chunk_len=25 \
+--num_chunk_probs='{"1":"0.5", "2":"0.125", "3":"0.0625", "4":"0.0625", "5":"0.0625", "6":"0.0625", "7":"0.0625", "8":"0.0625"}' \
+--max_val_samples_per_ds=100 \
+--seed=1 \
+--use_per_ctx_average_loss=True \
+--torch_empty_cache_steps=10

src/ctx_to_lora/data/definitions.py

@@ -10,6 +10,16 @@ CTX_AFFIXES = {
         "prefix": [2, 106, 1645, 110],  # <bos><start_of_turn>user\n\n\n
         "suffix": [107, 108, 106, 2516, 108],  # <end_of_turn>\n<start_of_turn>model\n
     },
+    "mistralai/Mistral-7B-Instruct-v0.2": {
+        "prefix": [1, 733, 16289, 28793, 28705, 13, 13],  # `<s> [INST] \n\n`
+        "suffix": [733, 28748, 16289, 28793],  # ` [/INST] `
+    },
+    "Qwen/Qwen3-4B-Instruct-2507": {
+        # `<|im_start|>system\n<|im_end|>\n<|im_start|>user\n`
+        "prefix": [151644, 8948, 198, 151645, 198, 151644, 872, 198],
+        # `<|im_end|>\n<|im_start|>assistant\n`
+        "suffix": [151645, 198, 151644, 77091, 198],
+    },
 }