230 lines
7.7 KiB
Python
230 lines
7.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Download ToolBench + APIGen-MT + ToolACE
|
|
and convert them to Qwen 2.5 SFT JSONL format
|
|
(with tool calling / function calling support).
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import gzip
|
|
import tarfile
|
|
import zipfile
|
|
import requests
|
|
from pathlib import Path
|
|
from tqdm import tqdm
|
|
from datasets import load_dataset
|
|
from huggingface_hub import hf_hub_download, snapshot_download
|
|
|
|
# ====================== CONFIG ======================
|
|
OUTPUT_DIR = Path("./qwen25_tool_sft")
|
|
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
FINAL_JSONL = OUTPUT_DIR / "tool_sft_qwen25.jsonl"
|
|
|
|
# ====================================================
|
|
|
|
def download_file(url: str, dest: Path):
|
|
if dest.exists():
|
|
print(f"[skip] {dest.name} already exists")
|
|
return
|
|
print(f"Downloading {url} ...")
|
|
with requests.get(url, stream=True) as r:
|
|
r.raise_for_status()
|
|
total = int(r.headers.get("content-length", 0))
|
|
with open(dest, "wb") as f, tqdm(total=total, unit="B", unit_scale=True) as pbar:
|
|
for chunk in r.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
pbar.update(len(chunk))
|
|
|
|
|
|
def to_qwen_messages(system: str | None, conversations: list[dict]) -> dict:
|
|
"""
|
|
Convert a list of turns into Qwen 2.5 messages format.
|
|
conversations: list of {"from": "human/gpt/function/...", "value": "..."}
|
|
"""
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
|
|
for turn in conversations:
|
|
role = turn.get("from", "").lower()
|
|
content = turn.get("value", "").strip()
|
|
if not content:
|
|
continue
|
|
|
|
if role in ("human", "user"):
|
|
messages.append({"role": "user", "content": content})
|
|
elif role in ("gpt", "assistant"):
|
|
messages.append({"role": "assistant", "content": content})
|
|
elif role in ("function", "tool", "observation"):
|
|
# Qwen-style tool response
|
|
messages.append({"role": "tool", "content": content})
|
|
else:
|
|
# fallback
|
|
messages.append({"role": "user", "content": content})
|
|
|
|
return {"messages": messages}
|
|
|
|
|
|
# ----------------------------------------------------
|
|
# 1. ToolBench (official)
|
|
# ----------------------------------------------------
|
|
def process_toolbench():
|
|
print("\n=== ToolBench ===")
|
|
# ToolBench is available on Hugging Face
|
|
try:
|
|
ds = load_dataset("ToolBench/ToolBench", split="train", trust_remote_code=True)
|
|
except Exception:
|
|
# fallback to the processed version that many people use
|
|
ds = load_dataset("lmsys/toolbench", split="train")
|
|
|
|
count = 0
|
|
with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
|
|
for sample in tqdm(ds, desc="ToolBench"):
|
|
# ToolBench usually has "conversations" or "messages"
|
|
convs = sample.get("conversations") or sample.get("messages") or []
|
|
if not convs:
|
|
continue
|
|
|
|
# Some versions already have role/content
|
|
if isinstance(convs[0], dict) and "role" in convs[0]:
|
|
messages = []
|
|
for m in convs:
|
|
role = m.get("role", "user")
|
|
content = m.get("content", "")
|
|
if role == "function":
|
|
role = "tool"
|
|
messages.append({"role": role, "content": content})
|
|
record = {"messages": messages}
|
|
else:
|
|
record = to_qwen_messages(None, convs)
|
|
|
|
if len(record["messages"]) >= 2:
|
|
fout.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
print(f"ToolBench → {count} samples")
|
|
|
|
|
|
# ----------------------------------------------------
|
|
# 2. APIGen-MT (multi-turn tool calling)
|
|
# ----------------------------------------------------
|
|
def process_apigen_mt():
|
|
print("\n=== APIGen-MT ===")
|
|
# Common locations / names
|
|
possible = [
|
|
"Salesforce/APIGen-MT",
|
|
"Salesforce/xLAM-APIGen",
|
|
"Salesforce/APIGen",
|
|
]
|
|
ds = None
|
|
for name in possible:
|
|
try:
|
|
ds = load_dataset(name, split="train")
|
|
print(f"Loaded {name}")
|
|
break
|
|
except Exception:
|
|
continue
|
|
|
|
if ds is None:
|
|
print("APIGen-MT not found on HF under common names. Skipping.")
|
|
return
|
|
|
|
count = 0
|
|
with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
|
|
for sample in tqdm(ds, desc="APIGen-MT"):
|
|
# APIGen usually has "messages" already close to OpenAI format
|
|
messages = sample.get("messages") or sample.get("conversations")
|
|
if not messages:
|
|
continue
|
|
|
|
# Normalize role names
|
|
normalized = []
|
|
for m in messages:
|
|
role = m.get("role", "user").lower()
|
|
content = m.get("content", "")
|
|
if role == "function":
|
|
role = "tool"
|
|
normalized.append({"role": role, "content": content})
|
|
|
|
if len(normalized) >= 2:
|
|
fout.write(json.dumps({"messages": normalized}, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
print(f"APIGen-MT → {count} samples")
|
|
|
|
|
|
# ----------------------------------------------------
|
|
# 3. ToolACE
|
|
# ----------------------------------------------------
|
|
def process_toolace():
|
|
print("\n=== ToolACE ===")
|
|
possible = [
|
|
"Team-ACE/ToolACE",
|
|
"ToolACE/ToolACE",
|
|
"microsoft/ToolACE",
|
|
]
|
|
ds = None
|
|
for name in possible:
|
|
try:
|
|
ds = load_dataset(name, split="train")
|
|
print(f"Loaded {name}")
|
|
break
|
|
except Exception:
|
|
continue
|
|
|
|
if ds is None:
|
|
print("ToolACE not found under common names. Trying alternative...")
|
|
# Some people host processed versions
|
|
try:
|
|
ds = load_dataset("json", data_files="https://huggingface.co/datasets/Team-ACE/ToolACE/resolve/main/data/train.json")
|
|
except Exception:
|
|
print("Could not load ToolACE. Skipping.")
|
|
return
|
|
|
|
count = 0
|
|
with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
|
|
for sample in tqdm(ds, desc="ToolACE"):
|
|
messages = sample.get("messages") or sample.get("conversations") or []
|
|
if not messages:
|
|
continue
|
|
|
|
normalized = []
|
|
for m in messages:
|
|
if isinstance(m, dict):
|
|
role = m.get("role", m.get("from", "user")).lower()
|
|
content = m.get("content", m.get("value", ""))
|
|
else:
|
|
continue
|
|
if role in ("function", "observation"):
|
|
role = "tool"
|
|
elif role in ("human", "user"):
|
|
role = "user"
|
|
elif role in ("gpt", "assistant"):
|
|
role = "assistant"
|
|
normalized.append({"role": role, "content": content})
|
|
|
|
if len(normalized) >= 2:
|
|
fout.write(json.dumps({"messages": normalized}, ensure_ascii=False) + "\n")
|
|
count += 1
|
|
print(f"ToolACE → {count} samples")
|
|
|
|
|
|
# ----------------------------------------------------
|
|
# Main
|
|
# ----------------------------------------------------
|
|
if __name__ == "__main__":
|
|
# Clear previous output if you want a fresh file
|
|
if FINAL_JSONL.exists():
|
|
print(f"Removing old {FINAL_JSONL}")
|
|
FINAL_JSONL.unlink()
|
|
|
|
process_toolbench()
|
|
process_apigen_mt()
|
|
process_toolace()
|
|
|
|
# Final stats
|
|
total = sum(1 for _ in open(FINAL_JSONL, "r", encoding="utf-8"))
|
|
print(f"\n✅ Done! Total samples written → {FINAL_JSONL}")
|
|
print(f" Total lines: {total}")
|
|
print("\nYou can now use this JSONL for Qwen2.5 SFT (tool calling / function calling).")
|