NVIDIA公开Open‑SWE‑Traces数据集助力打造高质量SFT训练集

2 阅读9分钟应用

关键步骤概览

  1. 数据流式获取:使用 datasets.load_dataset 从 HuggingFace 直接读取 nvidia/Open‑SWE‑Traces,避免本地完整下载。
  2. 轨迹标准化:将多轮对话统一为结构化 JSON,提取 rolecontent,并过滤异常格式。
  3. 补丁解析:统计每个 model_patch 的文件数、增删行数及文件后缀分布。
  4. 特征构建:计算消息数、角色分布、工具调用次数、整体令牌估算等指标,形成分析 DataFrame。
  5. 可视化分析:语言分布、分辨率率、工具使用频次、令牌长度等通过柱状图、散点图直观呈现。
  6. SFT 子集筛选:基于分辨率、令牌上限(默认 32k)、语言过滤与补丁可用性四大条件,生成符合训练需求的样本集合。
  7. 导出:将完整分析表保存为 CSV,SFT 数据保存为 JSONL,便于后续 datasets.load_dataset 直接加载。

数据流式获取与预处理

from datasets import load_dataset
DATASET = "nvidia/Open‑SWE‑Traces"
AGENTS = ["openhands", "sweagent"]
MODELS = ["minimax_m25", "qwen35_122b"]
PER_COMBO = 400  # 每组抽样数量

def stream_take(agent, model, n):
    ds = load_dataset(DATASET, agent, split=model, streaming=True)
    rows = []
    for ex in itertools.islice(ds, n):
        ex = dict(ex)
        ex["_agent"], ex["_model"] = agent, model
        rows.append(ex)
    return rows

通过循环组合 agent × model,一次性拉取约 1.6k 条轨迹,打印每批拉取量以确认网络连通性。整个过程在 Colab 中几分钟即可完成。

轨迹解析与特征提取

import json, re, pandas as pd
from collections import Counter

# 统一消息文本
def message_text(msg):
    if not isinstance(msg, dict):
        return ""
    content = msg.get("content", "")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "
".join(part.get("text", "") for part in content if isinstance(part, dict))
    return str(content)

# 统计角色出现次数
def role_counts(traj):
    cnt = Counter()
    for m in traj:
        cnt[m.get("role", "unknown")] += 1
    return cnt

# 解析 diff 补丁
def parse_patch(diff):
    if not diff or not isinstance(diff, str):
        return 0,0,0,[],Counter()
    files,add,del_ = [],0,0
    exts = Counter()
    for line in diff.splitlines():
        if line.startswith("diff --git"):
            parts = line.split()
            if len(parts) >= 3:
                path = parts[2][2:] if parts[2].startswith("a/") else parts[2]
                files.append(path)
                ext = path.split(".")[-1].lower()
                exts[ext] += 1
        elif line.startswith("+") and not line.startswith("+++"):
            add += 1
        elif line.startswith("-") and not line.startswith("---"):
            del_ += 1
    return len(files), add, del_, files, exts

对每条记录调用上述函数后,得到包含 instance_id、repo、language、resolved、n_messages、n_tool、patch_churn、traj_tokens 等字段的字典列表,随后 pd.DataFrame 合并为分析表。

统计分析与可视化

  • 语言分布:Python、JavaScript 等占比最高。
  • 分辨率率:已解决轨迹约 62%,不同语言的成功率差异显著。
  • 工具使用bash_blocksearch_file 等前 12 名工具占总调用的 78%。
  • 令牌预算:50% 轨迹在 2.4k 令牌以内,95% 以内不超过 12k,只有少数超出 32k 上限。

通过 plt.histplt.barh 可快速生成上述可视化,帮助判断哪些轨迹适配现有上下文窗口。

构建 SFT 子集

MAX_SFT_TOKENS = 32_000
SFT_REQUIRE_RESOLVED = True
SFT_LANGUAGES = None  # 如需限定语言可设为 ['python']

def passes_filters(rec, raw):
    if SFT_REQUIRE_RESOLVED and rec["resolved"] != 1:
        return False
    if rec["traj_tokens"] > MAX_SFT_TOKENS:
        return False
    if SFT_LANGUAGES and rec["language"] not in SFT_LANGUAGES:
        return False
    if not (raw.get("model_patch") or "").strip():
        return False
    return True

sft_examples = []
for rec, raw in zip(records, raw_rows):
    if not passes_filters(rec, raw):
        continue
    msgs = [{"role": m.get("role"), "content": message_text(m)} for m in normalize_trajectory(raw["trajectory"])]
    sft_examples.append({
        "instance_id": rec["instance_id"],
        "repo": rec["repo"],
        "language": rec["language"],
        "agent": rec["agent"],
        "model": rec["model"],
        "messages": msgs,
        "text": "
".join(f"<|im_start|>{m['role']}
{m['content'].strip()}<|im_end|>" for m in msgs),
        "model_patch": raw.get("model_patch"),
        "approx_tokens": rec["traj_tokens"]
    })

筛选后约保留 18% 的原始轨迹,形成可直接喂入 LLM 微调的 ChatML 格式文本。

导出与后续使用

pd.DataFrame(sft_examples).to_json('open_swe_sft.jsonl', orient='records', lines=True, force_ascii=False)

生成的 open_swe_sft.jsonlopen_swe_traces_analysis.csv 均可在 Colab 左侧文件面板下载,也可通过 datasets.load_dataset('json', data_files='open_swe_sft.jsonl') 直接加载到训练管线。后续可在更大规模抽样、语言特化或模型专属模板上继续扩展此框架。


通过本套完整工作流,研究者和工程师能够快速从公开的 agentic 软件工程轨迹中提炼出高质量的监督微调数据,为代码生成、错误修复以及工具调用等下游任务提供可靠的训练基准。

本文是对第三方新闻源的主观解读。消息可能出现过时、不准确、歧义或错误的地方,仅供参考使用。点击此处查看消息源。