From a074fd3305a34c07f8d587088c0295d81e2c690f Mon Sep 17 00:00:00 2001
From: 51616 <rujikorn.ch@gmail.com>
Date: Wed, 8 Oct 2025 05:17:50 +0000
Subject: [PATCH] self-gen data misaligned fixed + niah reproducible +
 vllm-tokenizers version combo

---
 README.md                   | 13 +++++---
 data/self_generate_qa.py    |  2 +-
 install.sh                  |  1 +
 pyproject.toml              |  3 +-
 scripts/niah/1-train.sh     | 62 ++++++++++++++++++++++++-------------
 scripts/niah/2-eval-test.sh |  1 +
 scripts/niah/README.md      |  8 +++++
 uv.lock                     | 36 +++------------------
 8 files changed, 64 insertions(+), 62 deletions(-)
 create mode 100755 scripts/niah/2-eval-test.sh
 create mode 100644 scripts/niah/README.md

diff --git a/README.md b/README.md
index b57cd52..0fad415 100644
--- a/README.md
+++ b/README.md
@@ -53,11 +53,14 @@ print(tokenizer.decode(outputs[0]))
 
 ### 🎮 Interactive Demo [WIP]
 ```bash
-uv run webui/demo.py
+uv run demo/app.py
 ```
 
 ### 🧪 Experimental Scripts
-| Experiment                           | Data prep                                                                                                                           | Training                                  | Evaluation                                                                                                                               | Notes                                                                                                                               |
-| ------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
-| [Main experiment](scripts/main_exp/) | `uv run bash scripts/main_exp/0-download_data.sh`<br>`uv run bash scripts/main_exp/gen_data.sh` *(optional, rebuilds from scratch)* | `uv run bash scripts/main_exp/1-train.sh` | `uv run bash scripts/main_exp/eval/<script>.sh` (pick from `base_model.sh`, `cd.sh`, `cd_oracle.sh`, `d2l.sh`, `llmlingua.sh`, `t2l.sh`) | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
-| [NIAH](scripts/niah/)                | `uv run bash scripts/niah/0-gen_data.sh`                                                                                            | `uv run bash scripts/niah/1-train.sh`     | `uv run bash scripts/niah/2-eval.sh`                                                                                                     | Run the scripts in order; data generation only needs to happen once                                                                 |
+To run any of the following scripts, use `uv run $PATH_TO_SCRIPT` from the root of this project.
+
+
+| Experiment                           | Data prep                             | Training                      | Evaluation                   | Notes                                                                                                                               |
+| ------------------------------------ | ------------------------------------- | ----------------------------- | ---------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- |
+| [Main experiment](scripts/main_exp/) | `scripts/main_exp/0-download_data.sh` | `scripts/main_exp/1-train.sh` | `scripts/main_exp/eval/*.sh` | Downloading data is fastest; regenerate only if you need fresh synthetic data. Evaluation scripts reproduce the main paper metrics. |
+| [NIAH](scripts/niah/)                | `scripts/niah/0-gen_data.sh`          | `scripts/niah/1-train.sh`     | `scripts/niah/2-eval.sh`     | Run the scripts in order; data generation only needs to happen once                                                                 |
diff --git a/data/self_generate_qa.py b/data/self_generate_qa.py
index aef4b91..67a75bf 100644
--- a/data/self_generate_qa.py
+++ b/data/self_generate_qa.py
@@ -233,7 +233,7 @@ def self_generate(
 
     tk = get_tokenizer(args.vllm_model, train=True)
 
-    self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"]
+    self_qa_intx_tokens = tk(SELF_QA_INTX, add_special_tokens=False)["input_ids"][1:]
     if args.remove_qa_template:
         self_qa_intx_tokens = tk("\n\n", add_special_tokens=False)["input_ids"]
     n_self_qa_intx_tokens = len(self_qa_intx_tokens)
diff --git a/install.sh b/install.sh
index 48afaef..8b1e128 100755
--- a/install.sh
+++ b/install.sh
@@ -7,6 +7,7 @@ uv self update
 uv venv --python 3.10 --seed
 uv pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --torch-backend=cu124
 uv sync
+uv pip install tokenizers==0.21.0
 uv pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 uv pip install flashinfer-python==0.2.2 -i https://flashinfer.ai/whl/cu124/torch2.6
 
diff --git a/pyproject.toml b/pyproject.toml
index 3b8b270..0cb77c7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,7 +30,7 @@ dependencies = [
     "pandas",
     "plotly",
     "rouge-score",
-    "vllm==0.8.4",
+    "vllm==0.8.5.post1",
     "huggingface-hub[hf-transfer]>=0.32.0",
     "opt-einsum>=3.4.0",
     "kagglehub[hf-datasets]>=0.3.12",
@@ -38,7 +38,6 @@ dependencies = [
     "bitsandbytes>=0.46.1",
     "google-cloud-storage>=3.2.0",
     "wonderwords>=2.2.0",
-    "tokenizers==0.21.0",
     "llmlingua>=0.2.2",
 ]
 
diff --git a/scripts/niah/1-train.sh b/scripts/niah/1-train.sh
index 97b4ad7..f28b261 100755
--- a/scripts/niah/1-train.sh
+++ b/scripts/niah/1-train.sh
@@ -1,22 +1,40 @@
-WANDB_MODE=disabled uv run train.py configs/toy_exp/ctx_magic_number_32_256.yaml \
---model_name_or_path=google/gemma-2-2b-it \
---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=1 --use_sequence_packing=True \
---max_packed_inp_len=2048 --max_packed_ctx_len=2048 \
---dataloader_num_workers=0 --dataloader_prefetch_factor=None \
---ctx_encoder_type=early_exit --n_latent_queries=208 \
---eval_on_start=True --lora_r=8 \
---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
+ WANDB_MODE=disabled run uv run train.py \
+ configs/niah_exp/ctx_magic_number_32_256.yaml \
+ --model_name_or_path=google/gemma-2-2b-it \
+ --num_train_epochs=1 \
+ --per_device_train_batch_size=-1 \
+ --gradient_accumulation_steps=16 \
+ --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=1.5 \
+ --use_sequence_packing=True \
+ --max_packed_inp_len=4096 \
+ --max_packed_ctx_len=4096 \
+ --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
\ No newline at end of file
diff --git a/scripts/niah/2-eval-test.sh b/scripts/niah/2-eval-test.sh
new file mode 100755
index 0000000..ef88394
--- /dev/null
+++ b/scripts/niah/2-eval-test.sh
@@ -0,0 +1 @@
+WANDB_MODE=disabled uv run run_eval.py --checkpoint_path CHECKPOINT_PATH --datasets ctx_magic_number_32_1024 ctx_magic_number_1024_2048 ctx_magic_number_3072_4096 ctx_magic_number_7168_8192 ctx_magic_number_15360_16384 ctx_magic_number_28672_32768 ctx_magic_number_57344_65536 ctx_magic_number_122880_131072 --max_ctx_chunk_len=1024 --split test --eval_batch_size_gen=4
\ No newline at end of file
diff --git a/scripts/niah/README.md b/scripts/niah/README.md
new file mode 100644
index 0000000..24403bc
--- /dev/null
+++ b/scripts/niah/README.md
@@ -0,0 +1,8 @@
+# NIAH experiment
+```bash
+# run the scripts in this order
+# data generation is only needed to be run once
+uv run bash scripts/niah/0-gen_data.sh
+uv run bash scripts/niah/1-train.sh
+uv run bash scripts/niah/2-eval.sh
+```
\ No newline at end of file
diff --git a/uv.lock b/uv.lock
index 4629f44..b687ded 100644
--- a/uv.lock
+++ b/uv.lock
@@ -856,7 +856,6 @@ dependencies = [
     { name = "setuptools" },
     { name = "tensorboard" },
     { name = "tensorboardx" },
-    { name = "tokenizers" },
     { name = "torchmetrics" },
     { name = "transformers" },
     { name = "vllm" },
@@ -894,10 +893,9 @@ requires-dist = [
     { name = "setuptools" },
     { name = "tensorboard" },
     { name = "tensorboardx" },
-    { name = "tokenizers", specifier = "==0.21.0" },
     { name = "torchmetrics" },
     { name = "transformers", specifier = "==4.51.3" },
-    { name = "vllm", specifier = "==0.8.4" },
+    { name = "vllm", specifier = "==0.8.5.post1" },
     { name = "wandb" },
     { name = "wonderwords", specifier = ">=2.2.0" },
 ]
@@ -5428,30 +5426,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e6/34/ebdc18bae6aa14fbee1a08b63c015c72b64868ff7dae68808ab500c492e2/tinycss2-1.4.0-py3-none-any.whl", hash = "sha256:3a49cf47b7675da0b15d0c6e1df8df4ebd96e9394bb905a5775adb0d884c5289", size = 26610, upload-time = "2024-10-24T14:58:28.029Z" },
 ]
 
-[[package]]
-name = "tokenizers"
-version = "0.21.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "huggingface-hub" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/20/41/c2be10975ca37f6ec40d7abd7e98a5213bb04f284b869c1a24e6504fd94d/tokenizers-0.21.0.tar.gz", hash = "sha256:ee0894bf311b75b0c03079f33859ae4b2334d675d4e93f5a4132e1eae2834fe4", size = 343021, upload-time = "2024-11-27T13:11:23.89Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b0/5c/8b09607b37e996dc47e70d6a7b6f4bdd4e4d5ab22fe49d7374565c7fefaf/tokenizers-0.21.0-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3c4c93eae637e7d2aaae3d376f06085164e1660f89304c0ab2b1d08a406636b2", size = 2647461, upload-time = "2024-11-27T13:11:07.911Z" },
-    { url = "https://files.pythonhosted.org/packages/22/7a/88e58bb297c22633ed1c9d16029316e5b5ac5ee44012164c2edede599a5e/tokenizers-0.21.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:f53ea537c925422a2e0e92a24cce96f6bc5046bbef24a1652a5edc8ba975f62e", size = 2563639, upload-time = "2024-11-27T13:11:05.908Z" },
-    { url = "https://files.pythonhosted.org/packages/f7/14/83429177c19364df27d22bc096d4c2e431e0ba43e56c525434f1f9b0fd00/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b177fb54c4702ef611de0c069d9169f0004233890e0c4c5bd5508ae05abf193", size = 2903304, upload-time = "2024-11-27T13:10:51.315Z" },
-    { url = "https://files.pythonhosted.org/packages/7e/db/3433eab42347e0dc5452d8fcc8da03f638c9accffefe5a7c78146666964a/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:6b43779a269f4629bebb114e19c3fca0223296ae9fea8bb9a7a6c6fb0657ff8e", size = 2804378, upload-time = "2024-11-27T13:10:53.513Z" },
-    { url = "https://files.pythonhosted.org/packages/57/8b/7da5e6f89736c2ade02816b4733983fca1c226b0c42980b1ae9dc8fcf5cc/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9aeb255802be90acfd363626753fda0064a8df06031012fe7d52fd9a905eb00e", size = 3095488, upload-time = "2024-11-27T13:11:00.662Z" },
-    { url = "https://files.pythonhosted.org/packages/4d/f6/5ed6711093dc2c04a4e03f6461798b12669bc5a17c8be7cce1240e0b5ce8/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d8b09dbeb7a8d73ee204a70f94fc06ea0f17dcf0844f16102b9f414f0b7463ba", size = 3121410, upload-time = "2024-11-27T13:10:55.674Z" },
-    { url = "https://files.pythonhosted.org/packages/81/42/07600892d48950c5e80505b81411044a2d969368cdc0d929b1c847bf6697/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:400832c0904f77ce87c40f1a8a27493071282f785724ae62144324f171377273", size = 3388821, upload-time = "2024-11-27T13:10:58.401Z" },
-    { url = "https://files.pythonhosted.org/packages/22/06/69d7ce374747edaf1695a4f61b83570d91cc8bbfc51ccfecf76f56ab4aac/tokenizers-0.21.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e84ca973b3a96894d1707e189c14a774b701596d579ffc7e69debfc036a61a04", size = 3008868, upload-time = "2024-11-27T13:11:03.734Z" },
-    { url = "https://files.pythonhosted.org/packages/c8/69/54a0aee4d576045b49a0eb8bffdc495634309c823bf886042e6f46b80058/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:eb7202d231b273c34ec67767378cd04c767e967fda12d4a9e36208a34e2f137e", size = 8975831, upload-time = "2024-11-27T13:11:10.32Z" },
-    { url = "https://files.pythonhosted.org/packages/f7/f3/b776061e4f3ebf2905ba1a25d90380aafd10c02d406437a8ba22d1724d76/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:089d56db6782a73a27fd8abf3ba21779f5b85d4a9f35e3b493c7bbcbbf0d539b", size = 8920746, upload-time = "2024-11-27T13:11:13.238Z" },
-    { url = "https://files.pythonhosted.org/packages/d8/ee/ce83d5ec8b6844ad4c3ecfe3333d58ecc1adc61f0878b323a15355bcab24/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:c87ca3dc48b9b1222d984b6b7490355a6fdb411a2d810f6f05977258400ddb74", size = 9161814, upload-time = "2024-11-27T13:11:16.675Z" },
-    { url = "https://files.pythonhosted.org/packages/18/07/3e88e65c0ed28fa93aa0c4d264988428eef3df2764c3126dc83e243cb36f/tokenizers-0.21.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:4145505a973116f91bc3ac45988a92e618a6f83eb458f49ea0790df94ee243ff", size = 9357138, upload-time = "2024-11-27T13:11:20.09Z" },
-    { url = "https://files.pythonhosted.org/packages/15/b0/dc4572ca61555fc482ebc933f26cb407c6aceb3dc19c301c68184f8cad03/tokenizers-0.21.0-cp39-abi3-win32.whl", hash = "sha256:eb1702c2f27d25d9dd5b389cc1f2f51813e99f8ca30d9e25348db6585a97e24a", size = 2202266, upload-time = "2024-11-27T13:11:28.784Z" },
-    { url = "https://files.pythonhosted.org/packages/44/69/d21eb253fa91622da25585d362a874fa4710be600f0ea9446d8d0217cec1/tokenizers-0.21.0-cp39-abi3-win_amd64.whl", hash = "sha256:87841da5a25a3a5f70c102de371db120f41873b854ba65e52bccd57df5a3780c", size = 2389192, upload-time = "2024-11-27T13:11:25.724Z" },
-]
 
 [[package]]
 name = "tomli"
@@ -5673,7 +5647,6 @@ dependencies = [
     { name = "regex" },
     { name = "requests" },
     { name = "safetensors" },
-    { name = "tokenizers" },
     { name = "tqdm" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/f1/11/7414d5bc07690002ce4d7553602107bf969af85144bbd02830f9fb471236/transformers-4.51.3.tar.gz", hash = "sha256:e292fcab3990c6defe6328f0f7d2004283ca81a7a07b2de9a46d67fd81ea1409", size = 8941266, upload-time = "2025-04-14T08:15:00.485Z" }
@@ -5877,7 +5850,7 @@ wheels = [
 
 [[package]]
 name = "vllm"
-version = "0.8.4"
+version = "0.8.5.post1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "aiohttp" },
@@ -5925,7 +5898,6 @@ dependencies = [
     { name = "setuptools", marker = "python_full_version >= '3.12'" },
     { name = "six", marker = "python_full_version >= '3.12'" },
     { name = "tiktoken" },
-    { name = "tokenizers" },
     { name = "torch" },
     { name = "torchaudio" },
     { name = "torchvision" },
@@ -5936,9 +5908,9 @@ dependencies = [
     { name = "xformers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "xgrammar", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/e6/d6/9d412cdaa92c3ab6250cef51217d37395b2aa372c6c14f90b1668adbbf63/vllm-0.8.4.tar.gz", hash = "sha256:522b13dd16c6c773dec0cb4c42ea591623d03ef94d16db8128ece2600017e6ac", size = 6667631, upload-time = "2025-04-15T00:30:10.147Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/68/a4/66914f97cd0b2e8b56d4ab9e1695415fbf38c3bc6f2dc5abf0487a32d2ef/vllm-0.8.5.post1.tar.gz", hash = "sha256:5e5be78ee00637de4ee29f75ce86edc6c224c05d9e58d067a511eb83c3afe32d", size = 7337621, upload-time = "2025-05-02T22:31:09.951Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/8e/cb/03dc1299e0456ff3d58a11f63682ef29aaf5b1bd7f21bfe0690d7ce6fc40/vllm-0.8.4-cp38-abi3-manylinux1_x86_64.whl", hash = "sha256:e346749ee8df48cdcd935d00a7fc123a1e17d9904b064401e74fc6ad73b8104a", size = 294098962, upload-time = "2025-04-15T00:30:02.293Z" },
+    { url = "https://files.pythonhosted.org/packages/47/07/e8822ca825f88aec61a7af6210b85cd0e0a4d6f757081d3282f32b1b0912/vllm-0.8.5.post1-cp38-abi3-manylinux1_x86_64.whl", hash = "sha256:b7f0dbd46f82aac6b2e7489c2bbf1c90fdc2e56d3f1fff44549951db40814d5a", size = 326421557, upload-time = "2025-05-02T22:31:02.23Z" },
 ]
 
 [[package]]