KRLabsOrg/tool-output-extraction-swebench

--- language: - en license: apache-2.0 size_categories: - 10K<n<100K task_categories: - text-generation - token-classification tags: - tool-output - code - swe-bench - distillation - agent - context-compression - context-pruning --- # Tool Output Extraction Dataset [**Paper**](https://huggingface.co/papers/2604.04979) | [**Code**](https://github.com/KRLabsOrg/squeez) Training data for [**squeez**](https://github.com/KRLabsOrg/squeez) — a small model that prunes verbose coding agent tool output to only the evidence the agent needs next. ## Task **Task-conditioned context pruning of a single tool observation for coding agents.** Given a focused extraction query and one verbose tool output, return the smallest verbatim evidence block(s) the agent should read next. The model copies lines from the tool output — it never rewrites, summarizes, or invents content. Every line in the target exists verbatim in the source. ## Dataset Summary | | Train | Dev | Test | Total | |---|---:|---:|---:|---:| | **Samples** | 10,508 | 240 | 618 | 11,366 | **Data sources:** | Source | Samples | Description | |--------|--------:|-------------| | SWE-bench real data | 9,205 | Real tool output executed on cloned Python repos | | Synthetic multi-ecosystem | 1,697 | LLM-generated tool output for JS, Rust, Go, Docker, etc. | | Synthetic negatives | 575 | Mismatched task+output pairs where nothing is relevant | ## Why This Dataset Exists LLM coding agents waste 80–95% of context tokens on irrelevant tool output. When an agent reads a 500-line file to find one function, or runs `pytest` to find a failing test, most of the output is noise. Squeez trains small models to compress this output before it enters the agent's context window. ### Why v3? The previous dataset (v2) had three problems: 1. **Line-number leakage.** The synthetic data generation pipeline showed numbered output to the teacher LLM, and the teacher leaked those numbers into the target annotations. 87% of synthetic response lines contained line number prefixes (`2: npm WARN...`) that did not exist in the raw tool output. This trained the model to hallucinate formatting. 2. **No canonical truth format.** The dataset stored XML-wrapped ChatML prompts as the ground truth. This tangled the benchmark representation with one specific model's training format, making cross-model evaluation fragile. 3. **Task drift.** The labeling task drifted from "context pruning for an agent" toward "recover one arbitrary teacher-selected subset." Queries were either too vague (full issue descriptions) or too literal (grep-able string searches), neither matching the actual product use case. v3 fixes all three by introducing a canonical span-based representation, regenerating all labels with a focus on verbatim extraction, and manually curating the test set. ## How The Data Was Created ### Step 1: Generate Tool Calls (SWE-bench) Starting from 2,294 instances in the [SWE-bench](https://github.com/princeton-nlp/SWE-bench) test split, we simulated what a coding agent would do for each issue: read source files, grep for symbols, run tests, check git history, install packages. We generated 3–7 tool calls per instance, weighted toward common agent actions (read_file 28%, grep 18%, test_output 8%, etc.). ### Step 2: Execute on Real Repos Every tool call was executed against the actual repository checked out at the base commit (before the fix). `git grep` ran real searches, `pytest` ran real tests, `pip install -e .` did real installs. This produces authentic tool output with real file paths, error messages, formatting, and noise — things an LLM cannot reliably generate from scratch. ### Step 3: Generate Focused Queries For each (issue, tool_output) pair, a teacher LLM generated a focused extraction query — a short, concrete request for evidence rather than a full issue description or a literal grep pattern. Good queries: *"Find the traceback block that explains the import failure"*, *"Find the diff hunk that changes CSV parsing"* Bad queries: *"Fix the bug"*, *"Find all lines containing 'raise AttributeError'"* ### Step 4: Label Gold Spans A teacher LLM selected gold spans — contiguous blocks of lines in the raw tool output that answer the extraction query. The teacher saw numbered output as a reference interface, but the canonical target maps line numbers back to the raw unnumbered text. Every target line must exist verbatim in the source. The canonical representation is a list of `{start_line, end_line}` spans over the raw tool output. XML wrappers, ChatML formatting, and line number prefixes only appear in derived training files, never in the ground truth. ### Step 5: Synthetic Multi-Ecosystem Data SWE-bench only covers Python repositories, but real coding agents work across all ecosystems. We generated synthetic data for 15 tool types that SWE-bench cannot provide: `npm_install`, `npm_build`, `tsc`, `eslint`, `cargo_build`, `go_build`, `mvn_gradle`, `make_cmake`, `docker_build`, `docker_logs`, `terraform`, `kubectl`, `pip_install`, `mypy_pyright`, and `curl`. **Two-pass generation:** 1. An LLM generates a realistic task description and tool output in XML markers. Each tool type has a config with scenarios (e.g., "peer dependency conflict", "missing native module") and seed examples guiding realistic formatting. 2. Given the numbered output, a second LLM call selects relevant line numbers as JSON. These are mapped back to raw unnumbered source lines, validated (must exist verbatim, must be non-empty, reasonable compression ratio), and stored as canonical spans. ### Step 6: Hard Negatives To teach the model when to output nothing, 575 samples use intentionally mismatched task+output pairs — e.g., a React authentication task paired with a Rust borrow-checker error. The correct answer is an empty extraction. ### Step 7: Quality Filtering and Assembly Samples with empty spans (where the teacher found nothing relevant) were capped at 10% per tool type. Overly broad annotations (>60% of lines selected) were reviewed. Train/dev/test splits were assembled: - **SWE-bench**: split by repository (test: xarray, flask; dev: requests; train: all others) - **Synthetic**: split per tool type (10% test, 5% dev, 85% train) - **Negatives**: capped at ~10% of positives per tool type in test ### Step 8: Test Set Curation The held-out test set was manually reviewed and 111 samples were excluded: - **Near-duplicate np.unicode_ errors (63)**: The xarray repo on NumPy 2.0 produces the same `AttributeError: np.unicode_ was removed` on every `import xarray`. These identical errors across 65 different xarray instances were deduplicated to 2 representative samples. - **Trivial tiny outputs (39)**: Samples with 1–2 line output (e.g., lint "All checks passed!", "Python 3.12.9", single-line curl errors). Nothing to filter — not a meaningful benchmark. - **Overly broad spans (5)**: Samples selecting >50% of a large output, or spanning the entire top half of a file. - **Wrong annotations (4)**: Mislabeled tool types, spans pointing to wrong content, or vague queries without task context. The exclusion list is tracked in `test_exclusions.json` with per-sample reasons. ## Formats The dataset ships in three parallel formats, all derived from the same canonical spans: ### Canonical (`canonical_train/dev/test.jsonl`) The source of truth. Model-agnostic, no XML, no formatting artifacts. ```json { "instance_id": "django__django-11270", "source": "swe", "tool_type": "read_file", "query": "Find the code block that validates the referer in CsrfViewMiddleware", "background_task": "Fix CSRF validation bug when referer URL contains port number...", "tool_output": "raw output exactly as shown to the agent", "gold_spans": [ {"start_line": 41, "end_line": 52} ], "is_irrelevant": false, "command": "django/middleware/csrf.py" } ``` - `gold_spans` reference 1-indexed line numbers in `tool_output` - `is_irrelevant: true` means no lines are relevant (hard negative) - `query` is the focused extraction request; `background_task` is the full issue for provenance ### Generative / Qwen (`train/dev/test.jsonl`) ChatML-formatted for SFT training with Qwen or similar models. ```json { "prompt": "<|im_start|>system You prune verbose tool output...<|im_end|> <|im_start|>user <query> Find the code block... </query> <tool_output> 1: class CsrfViewMiddleware: 2: def _check_referer(self, request): ... </tool_output><|im_end|> <|im_start|>assistant ", "response": "<relevant_lines> 41: referer = request.META.get('HTTP_REFERER') 42: if referer is None: ... </relevant_lines>", "metadata": { "instance_id": "django__django-11270", "tool_type": "read_file", "source": "swe", "num_total_lines": 84, "num_relevant_lines": 12, "compression_ratio": 0.857 } } ``` ### Encoder (`encoder_train/dev/test.jsonl`) For token/line classification models (mmBERT, etc.). ## Tool Types 27 tool types across multiple ecosystems: | Ecosystem | Tool types | Source | |-----------|-----------|-------| | **Python** | read_file, grep, python, test_output, type_check, coverage, lint_output, build_output | SWE-bench | | **Python** | pip_install, curl | SWE-bench + synthetic | | **Git** | git_log, git_diff, git_blame, ls | SWE-bench | | **JavaScript/TypeScript** | npm_install, npm_build, tsc, eslint | Synthetic | | **Rust** | cargo_build | Synthetic | | **Go** | go_build | Synthetic | | **Java** | mvn_gradle | Synthetic | | **C/C++** | make_cmake | Synthetic | | **Infrastructure** | docker_build, docker_logs, terraform, kubectl | Synthetic | | **Python (type checking)** | mypy_pyright | Synthetic | ## Splits **SWE-bench data** is split by repository (zero instance overlap): - **Test**: `pydata/xarray`, `pallets/flask` - **Dev**: `psf/requests` - **Train**: all others (django, sympy, scikit-learn, sphinx, matplotlib, pytest, astropy, pylint, seaborn) **Synthetic data** is split per tool type: 10% test, 5% dev, 85% train. Hard negatives are capped at ~10% per tool type in test. ## Usage ```python from datasets import load_dataset ds = load_dataset("KRLabsOrg/tool-output-extraction-swebench") # Generative training splits print(ds) # DatasetDict({ # train: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 10508}) # dev: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 240}) # test: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 618}) # }) ``` ## Citation ```bibtex @misc{kovács2026squeeztaskconditionedtooloutputpruning, title={Squeez: Task-Conditioned Tool-Output Pruning for Coding Agents}, author={Ádám Kovács}, year={2026}, eprint={2604.04979}, archivePrefix={arXiv}, primaryClass={cs.SE}, url={https://arxiv.org/abs/2604.04979}, } ```language: - 英语 license: Apache-2.0 size_categories: - 10K<n<100K task_categories: - 文本生成 - Token分类 tags: - 工具输出（tool-output） - 代码（code） - SWE-bench - 知识蒸馏（distillation） - 智能体（agent） - 上下文压缩（context-compression） - 上下文剪枝（context-pruning） # 工具输出提取数据集 [**Paper**](https://huggingface.co/papers/2604.04979) | [**Code**](https://github.com/KRLabsOrg/squeez) 本数据集为[squeez](https://github.com/KRLabsOrg/squeez)的训练数据，该模型是一款轻量化模型，可将冗长的编码智能体工具输出剪枝至智能体后续所需的唯一证据片段。 ## 任务 **面向编码智能体的单工具观测结果任务感知上下文剪枝（task-conditioned context pruning）** 给定一个聚焦的提取查询与一段冗长的工具输出，返回智能体下一步所需读取的最简短逐字证据块。该模型仅从工具输出中复制原文行，绝不进行重写、摘要或内容编造，目标中的每一行均完整存在于源数据中。 ## 数据集概览 | | 训练集 | 验证集 | 测试集 | 总计 | |---|---:|---:|---:|---:| | **样本数** | 10508 | 240 | 618 | 11366 | **数据来源：** | 数据来源 | 样本数 | 描述 | |--------|--------:|-------------| | SWE-bench 真实数据 | 9205 | 在克隆的Python代码仓库上执行得到的真实工具输出 | | 多生态合成数据 | 1697 | 针对JavaScript、Rust、Go、Docker等开发生态生成的大语言模型工具输出 | | 合成负样本 | 575 | 任务与输出不匹配且无相关内容的样本对 | ## 本数据集的构建初衷 大语言模型编码智能体有80%~95%的上下文令牌被浪费在无关的工具输出上。当智能体读取一个500行的文件以查找某个函数，或运行`pytest`以定位失败的测试时，绝大多数输出均为噪声。Squeez旨在训练轻量化模型，在工具输出进入智能体上下文窗口前对其进行压缩。 ### 为何推出v3版本？ 此前的数据集（v2）存在三个问题： 1. **行号泄露**：合成数据生成流程会向教师大语言模型展示带行号的输出，导致教师模型将行号泄露至目标标注中。87%的合成响应行包含原始工具输出中不存在的行号前缀（如`2: npm WARN...`），这会使模型学习到编造格式的行为。 2. **无标准真值格式**：该数据集曾将XML封装的ChatML提示词作为真值存储，将基准测试的表示方式与特定模型的训练格式绑定，导致跨模型评估变得不可靠。 3. **任务偏移**：标注任务从“为智能体进行上下文剪枝”偏离为“恢复教师模型任意选择的子集”。查询语句要么过于模糊（如完整的问题描述），要么过于字面化（如可通过grep匹配的字符串搜索），均与实际产品的使用场景不符。 v3版本通过引入基于标准跨度（span）的表示方式、以逐字提取为核心重新生成所有标注，以及手动筛选与整理测试集来解决上述所有问题。 ## 数据构建流程 ### 步骤1：生成工具调用（SWE-bench） 我们从[SWE-bench](https://github.com/princeton-nlp/SWE-bench)测试划分中的2294个样本出发，模拟编码智能体针对每个问题的操作流程：读取源代码文件、grep搜索符号、运行测试、查看Git历史、安装依赖包。每个样本生成3~7个工具调用，权重向常见智能体操作倾斜（read_file占28%、grep占18%、test_output占8%等）。 ### 步骤2：在真实仓库中执行 每个工具调用均在基于基准提交（修复前）检出的真实仓库中执行。`git grep`执行真实搜索、`pytest`运行真实测试、`pip install -e .`执行真实安装。由此生成的工具输出包含真实文件路径、错误信息、格式与噪声，这些内容无法由大语言模型从零开始可靠生成。 ### 步骤3：生成聚焦查询 针对每个（问题，工具输出）样本对，教师大语言模型会生成一个聚焦的提取查询：一段简短且具体的证据请求，而非完整的问题描述或字面化的grep搜索模式。 优质查询示例：*"查找可解释导入失败的回溯块"*、*"查找修改CSV解析逻辑的diff块"* 劣质查询示例：*"修复该漏洞"*、*"查找所有包含'raise AttributeError'的行"* ### 步骤4：标注黄金跨度（gold spans） 教师大语言模型会选择黄金跨度：即原始工具输出中能够回答提取查询的连续行块。教师模型会将带行号的输出作为参考界面，但标准真值会将行号映射回原始无行号的文本。目标中的每一行必须完整存在于源数据中。 标准表示形式为原始工具输出上的`{start_line, end_line}`跨度列表。XML封装、ChatML格式与行号前缀仅会出现在衍生的训练文件中，不会出现在真值标注中。 ### 步骤5：多生态合成数据 SWE-bench仅覆盖Python仓库，但真实编码智能体需要适配所有开发生态。我们为SWE-bench无法提供的15种工具类型生成了合成数据：`npm_install`、`npm_build`、`tsc`、`eslint`、`cargo_build`、`go_build`、`mvn_gradle`、`make_cmake`、`docker_build`、`docker_logs`、`terraform`、`kubectl`、`pip_install`、`mypy_pyright`与`curl`。 **两步生成流程：** 1. 大语言先生成带有XML标记的真实任务描述与工具输出，每种工具类型都配有包含场景（如“对等依赖冲突”“缺失原生模块”）与示例的配置文件，以引导生成符合真实格式的内容。 2. 基于带行号的输出，调用第二个大语言模型以JSON格式选择相关行号，随后将其映射回原始无行号的源行，经过验证（必须完整存在于源数据中、非空且压缩比例合理）后存储为标准跨度格式。 ### 步骤6：难负样本 为了教会模型在无相关内容时输出空结果，我们构建了575个故意不匹配的任务与输出样本对，例如将React认证任务与Rust借用检查器错误输出配对，此时正确的提取结果为空。 ### 步骤7：质量过滤与数据集组装 每个工具类型的空跨度样本（教师模型未找到相关内容）占比被限制在10%以内。标注范围过宽（选择超过60%的行）的样本会被重新审核。训练集/验证集/测试集的划分规则如下： - **SWE-bench数据**：按仓库划分（测试集：xarray、flask；验证集：requests；训练集：其余所有仓库） - **合成数据**：按工具类型划分（10%测试集、5%验证集、85%训练集） - **负样本**：在测试集中每个工具类型的负样本占比被限制在约10%的正样本范围内。 ### 步骤8：测试集筛选与整理 我们对预留测试集进行了人工审核，共剔除111个样本，具体分类如下： - **近乎重复的np.unicode_错误（63个）**：NumPy 2.0版本下的xarray仓库会在每次`import xarray`时生成相同的`AttributeError: np.unicode_ was removed`错误。65个xarray样本中的此类重复错误被去重为2个代表性样本。 - **过于简单的短输出（39个）**：输出仅为1~2行的样本（如“所有检查通过！”的lint结果、“Python 3.12.9”版本信息、单行curl错误），此类样本无需剪枝，不具备基准测试意义。 - **标注范围过宽（5个）**：选择超过50%的大型输出内容，或覆盖文件前半部分全部内容的样本。 - **标注错误（4个）**：工具类型标注错误、跨度指向错误内容，或缺乏任务上下文的模糊查询。 样本剔除列表存储于`test_exclusions.json`文件中，包含每个样本的剔除原因。 ## 数据格式 本数据集提供三种并行格式，均基于同一套标准跨度生成： ### 标准格式（`canonical_train/dev/test.jsonl`） 作为所有数据的真值来源，该格式与模型无关，不包含XML与格式伪影。 json { "instance_id": "django__django-11270", "source": "swe", "tool_type": "read_file", "query": "Find the code block that validates the referer in CsrfViewMiddleware", "background_task": "Fix CSRF validation bug when referer URL contains port number...", "tool_output": "raw output exactly as shown to the agent", "gold_spans": [ {"start_line": 41, "end_line": 52} ], "is_irrelevant": false, "command": "django/middleware/csrf.py" } - `gold_spans` 引用`tool_output`中的1索引行号 - `is_irrelevant: true` 表示无相关行（难负样本） - `query` 为聚焦的提取请求；`background_task` 为完整的问题描述，用于溯源 ### 生成式/Qwen格式（`train/dev/test.jsonl`） 适配Qwen或同类模型的监督微调（SFT）训练，采用ChatML格式。 json { "prompt": "<|im_start|>system You prune verbose tool output...<|im_end|> <|im_start|>user <query> Find the code block... </query> <tool_output> 1: class CsrfViewMiddleware: 2: def _check_referer(self, request): ... </tool_output><|im_end|> <|im_start|>assistant ", "response": "<relevant_lines> 41: referer = request.META.get('HTTP_REFERER') 42: if referer is None: ... </relevant_lines>", "metadata": { "instance_id": "django__django-11270", "tool_type": "read_file", "source": "swe", "num_total_lines": 84, "num_relevant_lines": 12, "compression_ratio": 0.857 } } ### 编码器格式（`encoder_train/dev/test.jsonl`） 适用于Token/行分类模型（如mmBERT等）。 ## 工具类型 本数据集包含覆盖多个开发生态的27种工具类型： | 开发生态 | 工具类型 | 数据来源 | |-----------|-----------|-------| | **Python** | read_file, grep, python, test_output, type_check, coverage, lint_output, build_output | SWE-bench | | **Python** | pip_install, curl | SWE-bench + 合成数据 | | **Git** | git_log, git_diff, git_blame, ls | SWE-bench | | **JavaScript/TypeScript** | npm_install, npm_build, tsc, eslint | 合成数据 | | **Rust** | cargo_build | 合成数据 | | **Go** | go_build | 合成数据 | | **Java** | mvn_gradle | 合成数据 | | **C/C++** | make_cmake | 合成数据 | | **基础设施** | docker_build, docker_logs, terraform, kubectl | 合成数据 | | **Python（类型检查）** | mypy_pyright | 合成数据 | ## 数据集划分 **SWE-bench数据** 按仓库划分（样本无重叠）： - **测试集**：`pydata/xarray`、`pallets/flask` - **验证集**：`psf/requests` - **训练集**：其余所有仓库（django、sympy、scikit-learn、sphinx、matplotlib、pytest、astropy、pylint、seaborn） **合成数据** 按工具类型划分：10%作为测试集、5%作为验证集、85%作为训练集。测试集中每个工具类型的负样本占比被限制在约10%的正样本范围内。 ## 使用方法 python from datasets import load_dataset ds = load_dataset("KRLabsOrg/tool-output-extraction-swebench") # 生成式训练划分 print(ds) # DatasetDict({ # train: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 10508}) # dev: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 240}) # test: Dataset({features: ['prompt', 'response', 'metadata'], num_rows: 618}) # }) ## 引用 bibtex @misc{kovács2026squeeztaskconditionedtooloutputpruning, title={Squeez: Task-Conditioned Tool-Output Pruning for Coding Agents}, author={Ádám Kovács}, year={2026}, eprint={2604.04979}, archivePrefix={arXiv}, primaryClass={cs.SE}, url={https://arxiv.org/abs/2604.04979}, }