speedbench_64k128k

## 0. 现状回顾（已核实） - HF `nvidia/SPEED-Bench` 现有 config：`qualitative, throughput_1k/2k/8k/16k/32k`，每个 throughput config **1536 条 = high_entropy 512 + mixed 512 + low_entropy 512**。 - 每个 config 的目标 token 数用 `tiktoken` `o200k_base` 精确对齐（1k=1000 … 32k=32000）。 - parquet 中 turns 分两种： - **占位符**（`FULL BENCHMARK DATA SHOULD BE FETCHED FROM THE SOURCE USING SPECDEC_BENCH`）→ `prepare_data.py` 运行时从源拉取解析：HLE、AdaLEval StackSelect、AdaLEval TextSort、BAMBOO。 - **预解析**（直接在 parquet 里）：Gutenberg、WritingBench、RepoBench Python/Java、LCA。 - 每类 512 条的配比随长度桶变化：自然长度够的 source 占比上升，不够的下沉/消失。 **64k/128k 源数据可用性（已实测）：** | source | 64k | 128k | 说明 | |---|---|---|---| | AdaLEval StackSelect | ✅ `stackselect_64k.json`（200 条） | ✅ `stackselect_128k.json`（200 条） | opencompass.openxlab.space 已有 | | AdaLEval TextSort | ✅ `textsort_64k.json`（200 条） | ✅ `textsort_128k.json`（200 条） | 同上，自带 `prompt`/`num_tokens` 字段 | | HLE | ✅ few-shot 演示无限拼接 | ✅ 同上 | train 演示按类别采样 | | BAMBOO | ✅ 拼接 4~5 条 16k 对话 | ✅ 拼接 8~9 条 | 池子：showspred 100 条 + meetingpred 68 条 | | Gutenberg | ✅ 书长足够（32k 全部为截断） | ⚠️ 需从 ~2900 本书单筛更长 | 书单 NVIDIA 未公开，需重建 | | RepoBench Py/Java | ⚠️ 部分仓库可达 | ❌ 大部分不可达 | 受仓库大小限制，需实测 | | LCA | ⚠️ huge_context 部分可达 | ❌ 大部分不可达 | 受仓库大小限制，需实测 | | WritingBench | ❌ 太短 | ❌ | 只出现在 1k/2k，64k/128k 不参与 | --- ## 1. 总体流程 ``` 构造 64k/128k ├─ 1) 扩展 speed.py：config_type 增加 throughput_64k / throughput_128k │ _get_num_tokens_from_config 自动支持（regex 已通用） ├─ 2) 对"占位符源"：构造 parquet 时 source/src_id 指向 64k/128k 源文件 │ （StackSelect/TextSort 直接用已有 JSON；HLE 复用现有逻辑；BAMBOO 扩展拼接） ├─ 3) 对"预解析源"：重建 NVIDIA 模板 → 直接写入 parquet（Gutenberg/RepoBench/LCA） ├─ 4) 类别平衡：每类凑齐 512 条（low_entropy 为最大风险，见 §5） └─ 5) 产出 parquet + 校验：解析后每条 turns 精确 == 64000 / 128000 tokens ``` --- ## 2. 各 source 模板重建方案 ### 2.1 HLE（mixed，占位符，复用现有代码 ✅） - 现有 `_generate_hle_prompt`（`speed.py:448`）逻辑：同类别 train 演示拼接 + 目标问题，然后截断到 `num_tokens`。 - 64k/128k 只需把 `num_tokens` 传 64000/128000，**无需改代码**。 - 说明：演示从 `hle_train`（image=="" 过滤后）按类别采样，随机种子 `np.random.default_rng(42)` 固定；64k/128k 拼接更多演示即可。 ### 2.2 AdaLEval StackSelect（mixed，占位符，复用现有代码 ✅） - 源文件已存在：`http://opencompass.openxlab.space/utils/AdaLEval/stackselect_64k.json` / `_128k.json`（各 200 条，字段 `question_id/question/all_answers/answer/tags`）。 - 现有 `_generate_stackselect_prompt` + `_pad_or_truncate_prompt` 直接可用（已实测 32k 精确对齐 32000）。 - parquet 中 source 字段指向上述 URL，src_id = `question_id`。 ### 2.3 AdaLEval TextSort（low_entropy，占位符，复用现有代码 ✅） - 源文件已存在：`textsort_64k.json` / `_128k.json`（各 200 条，字段含 `prompt`，`num_tokens`≈63557/127357）。 - 现有 `_generate_textsort_prompt`（替换指令）+ `_pad_or_truncate_prompt` 直接可用（已实测 32k 精确对齐 32000）。 - parquet 中 source 指向上述 URL，src_id 格式 `textsort_<idx>`（同现有）。 ### 2.4 BAMBOO（high_entropy，占位符，需小改 ✅） - 源：GitHub `RUCAIBox/BAMBOO` 的 `showspred_16k.jsonl` / `meetingpred_16k.jsonl`（每条 content ~4k/~14k tokens）。 - 64k ≈ 拼接 4~5 条 meetingpred 或 ~16 条 showspred；128k ≈ 9 条 / 32 条。 - 现有 `_generate_bamboo_prompt` 已支持多 content 拼接（`START CONTENT N … END CONTENT`），只需： - `_fetch_all_turns_data` 中 BAMBOO 分支：`src_id` 解析为**多个**索引（现有 regex `_(\d+)` 已支持 `showspred_22_showspred_23` 这种多段格式）； - 选择"拼接多少条"的策略：按目标 token 数估算（每条 content 的 token 数），拼到 ≥ 目标附近，交给 `_pad_or_truncate_prompt` 收尾。 - 注意：现有代码在 BAMBOO 分支用 `num_tokens` 计算"answers_to_add_stop"逻辑的是 StackSelect，BAMBOO 分支走 `_generate_bamboo_prompt`，需确认 32k 的 2 条拼接是否由 NVIDIA 在 parquet 里预置（`showspred_22_showspred_23` 说明是），解析端只需把 src_ids 解析出来 select 即可 —— 现有代码已支持。**实测 2 条拼接 → 32000 精确对齐通过。** ### 2.5 Gutenberg（high_entropy，预解析，需重建 ⚠️） - 模板（已从 parquet 精确提取）： ``` "\nI want you to act as a book author. \nGiven a book, your objectives is to continue the book. Be coherent with the previous book and be creative in your response. Please continue the book in the same style as the book.\nThe content of the book is given below.\n\n" + <书全文(含 Gutenberg 头部)> + 截断到目标 ``` - 32k 的 486 条**全部是截断**（书都长于 32k）→ 64k/128k 同样按截断处理。 - **关键难点：书单**。src_id 形如 `gutenberg_open_books_<idx>`（idx 1~2908+），书单 NVIDIA 未公开。 - 重建方案（可选其一）： - (a) 用公开 Gutenberg 语料（如 `common-pile/project_gutenberg` 或 PG 官网文本）建英文书库 → 按"书全文 token ≥ 目标 + 阈值"过滤 → 顺序编号，src_id 格式照旧； - (b) 若贵方有 NVIDIA 内部书单可直接复用。 - 说明：模板与构造方法 100% 一致，书单近似（满足"尽量一致"）。 ### 2.6 RepoBench Python / Java（low_entropy，预解析，需重建 ⚠️） - 源：HF `tianyang/repobench_python_v1.1` / `repobench_java_v1.1`（revision 8a7cf0c8… / 0c2b0db4…），src_id 形如 `tianyang/repobench_python_v1.1-cross_file_random_298`（task 类型：cross_file_random / cross_file_first / in_file）。 - 模板（已提取）： ``` "\nYou are an AI assistant. Your job is to complete the code in the file below using the information provided about other files in the same repository in the context below. \nPlease complete the code in the file below, and provide a short explanation of the code.\n\nExample input format:\n\n# Repo Name: <repo_name>\n# Path: <repo_file_path>\n<Code snippet>\n\n....\n\n# Path: <repo_file_path>\n<cropped code snippet>\n\n\nExample output format:\n\n```\n<code completion>\n```\nExplanation:\n<explanation of the code completion>\n\n" + "# Repo Name: <repo>\n# Path: <file>\n" + <代码> + 填充到目标 ``` - 构造：取 dataset 行的 repo_name / file_path / context(identifier,path,snippet) / import_statement，按模板拼装后 `_pad_or_truncate_prompt`。 - 需实测：多少 repo 在拼接后能 ≥ 64k/128k（决定可入选数量）。 ### 2.7 LCA（low_entropy，预解析，需重建 ⚠️） - 源：HF `JetBrains-Research/lca-project-level-code-completion`（revision 1cdf6233…），config：huge_context(296)/large(270)/medium(224)/small(144)。 - 模板（已提取）： ``` "\nPlease fill in the missing code in the file using only the information provided in the context below. \nFor every [COMPLETE] marker in the file, output exactly _ONE_ line of code as the replacement. \nPlease output each line of code in a new line in the order they appear in the file, along with an explanation of the code.\n\nExample:\n```completion for line 1```, explanation\n```completion for line 2```, explanation\n" + 上下文文件们: "FILE NAME: <path>\n\nSTART FILE CONTENT\n\n<content>\n\nEND FILE CONTENT\n\n\n" + 目标文件: "COMPLETION FILE NAME: <path>\n\nSTART FILE CONTENT\n\n<含 [COMPLETE] 标记的代码>\n\nEND FILE CONTENT\n\n" ``` - 观察：32k 下 END FILE CONTENT 保留在末尾（先裁上下文、保目标文件完整性）；1k 下自然长度 > 目标也不截断。 - 重建：从源取 context 文件 + completion 文件，按模板拼装；长度裁剪策略：优先截上下文，保 completion 文件；过长仍超则保留原样（对齐 1k 行为）。 - 需实测：huge_context 中多少可达 64k/128k。 ### 2.8 WritingBench（不参与 64k/128k ❌） - query 最长 ~2k tokens，64k/128k 无样本。模板无需重建。 --- ## 3. 类别平衡策略（目标 512 / 512 / 512） | 类别 | source 贡献 | 64k 预计 | 128k 预计 | |---|---|---|---| | high_entropy | Gutenberg（书单筛长书）+ BAMBOO（拼接） | ~480 + ~32 | ~430 + ~82 | | mixed | HLE（few-shot）+ StackSelect（200 条源） | ~312 + 200 | ~312 + 200 | | low_entropy | TextSort(200) + RepoBench + LCA | 200 + 需实测 | 200 + 需实测（**最大风险**） | - high_entropy / mixed 可行性高；**low_entropy 若 RepoBench/LCA 可达样本不足，512 可能凑不齐**。 - 备选缓解： - TextSort 源文件只有 200 条——检查是否可用 `textsort_32k.json` 同款书籍追加构造更多 64k/128k 排序样本（书本身长度够，可自建 num_tokens 对齐版本）； - 或接受 low_entropy < 512（偏离原构造，需你确认）； - 或从 RepoBench 大仓库（cross_file）多选。 --- ## 4. 代码变更清单 ``` examples/specdec_bench/specdec_bench/datasets/speed.py ├─ config_type 增加 "throughput_64k", "throughput_128k" （_get_num_tokens_from_config 自动生效） ├─ DATASETS_AND_LOADERS_FUNCTIONS: 无新增源（全部复用现有 loader；BAMBOO 源 URL 用 16k 版） ├─ _fetch_all_turns_data: │ ├─ BAMBOO 分支：确认多 src_id 拼接 + 按目标条数选择（64k/128k 需拼更多条） │ ├─ HLE / StackSelect / TextSort 分支：现有逻辑直接支持（URL 换 64k/128k 文件） │ └─ （可选）新增 Gutenberg/RepoBench/LCA 的"本地已解析"处理：非占位符直接透传，无需代码 ├─ prepare_data.py: --config 的 choices 自动带上新 config（get_args(config_type)），无需改 └─ （可选）新增 build_throughput_64k_128k.py：生成占位符 parquet（见下） ``` **新增脚本 `build_speedbench_64k128k.py`（构造占位符 parquet）**： 1. 加载 StackSelect/TextSort 64k、128k 源 JSON → 生成占位符行（source=URL, src_id, category, sub_category, multiturn=False, difficulty=None, turns=[占位符]）。 2. HLE：从 hle_test 抽样（与 32k 同源、revision 021a3d71…）→ 占位符行。 3. BAMBOO：从 showspred/meetingpred_16k 选样本、决定拼接条数 → 占位符行（src_id 多段）。 4. Gutenberg/RepoBench/LCA：重建模板 → 直接写入**已解析 turns**（对齐 64k/128k）。 5. 平衡 512/512/512 → 写 parquet（字段与现有完全一致）。 **验证**：跑 `prepare_data.py --config throughput_64k`（HLE 下载较大）→ 检查每条 turns token 数 == 64000。 --- ## 5. 风险与未决问题 1. **low_entropy 数量**（最大风险）：RepoBench/LCA 64k/128k 可达样本数需实测（LCA huge_context 296 条、RepoBench 大 repo）；若不足 512 需决定缓解方案（§3）。 2. **Gutenberg 书单**：NVIDIA 未公开 `gutenberg_open_books_*` 书单（idx 1~2908）。用公开 PG 语料重建会"方法一致、书目近似"。需你确认是否接受，或是否有内部书单。 3. **BAMBOO 拼接策略**：拼接条数选择逻辑（32k 的 2 条是 NVIDIA 预置的）——64k/128k 按 token 估算拼接，与原文会有细微差异。 4. **HLE 采样随机性**：`_generate_hle_prompt` 用固定 rng(42)，但演示采样顺序依赖 hle_train 的行序 → 与 NVIDIA 的精确样本不会逐字相同（同一方法、不同演示组合）。可接受则 OK。 5. **数据集版本**：本方案基于 HF 当前版本（rev 487aa718，2026-02 更新）的数据结构；若后续 NVIDIA 官方更新了 speed.py（如已有 agentic config 分支），以官方为准。 --- ## 6. 交付物（✅ 已完成，2026-08-18） - ✅ `specdec_bench/datasets/speed.py`： - `config_type` 增加 `throughput_64k` / `throughput_128k` - BAMBOO URL 加载修复（datasets 5.x fsspec glob 回归 → 下载到本地缓存） - `_generate_hle_prompt` 精确收口改进（±1 边界漂移邻域扫描） - ✅ `build_speedbench_64k128k.py`（构建脚本，已放入 examples/specdec_bench/） - ✅ `/data1/xii/data_speed_64k128k/throughput_{64k,128k}/test.parquet`（占位符版，含 source/src_id 元数据） - ✅ `/data1/xii/data_speed_64k128k/final/throughput_{64k,128k}/test.parquet`（**完全解析版，可直接 `--dataset_path` 运行**） - ✅ 端到端验证：占位符版经 SPEEDBench `_prepare_mode=True` 全量解析成功 - 64k：1536 行，512/512/512，token 对齐 1519/1536 精确（17 条 ±1：16 条 HLE +1、1 条 repobench_java -1） - 128k：1536 行，512/512/512，token 对齐 1515/1536 精确（21 条 ±1，全部 HLE +1，源自 `_generate_hle_prompt` 固有边界漂移，NVIDIA 数据同样存在 ±1） ### 各 config 实际构成 | config | high_entropy (512) | mixed (512) | low_entropy (512) | |---|---|---|---| | 64k | Gutenberg 486 + BAMBOO 26 | HLE 412 + StackSelect 100 | TextSort 200 + LCA 100 + RepoBenchPy 62 + RepoBenchJava 150 | | 128k | Gutenberg 486 + BAMBOO 26 | HLE 412 + StackSelect 100 | TextSort 200 + LCA 312（huge 296 + large 16），RepoBench 128k 不可达（0） | ### 使用方式 ```bash # 直接跑（用已解析版） python run.py --dataset speed --dataset_path /data1/xii/data_speed_64k128k/final/throughput_64k ... # 或从占位符版重新解析（与官方 prepare_data.py 相同的解析逻辑） python -c " from specdec_bench.datasets.speed import SPEEDBench sb = SPEEDBench(config_name='data_speed_64k128k/throughput_64k', _prepare_mode=True) sb._resolved_dataset.to_parquet('out/test.parquet')" ``` ### 遗留说明 - HLE 解析的 ±1 偏差来自 NVIDIA `_generate_hle_prompt` 的 token 边界合并（已做改进但仍无法 100% 消除），与 NVIDIA 自身 32k 数据的 ±1 容差一致，对吞吐基准无影响。 - Gutenberg 书单为公开 PG 语料（common-pile/project_gutenberg）重建，书目与 NVIDIA 内部书单近似（方法一致、书目不同）。 - 原 `prepare_data.py` 的 `--config` 参数无法直接传本地路径（argparse choices 限制），解析需用上述等价调用。