get_tool_call.py
7.7 KB · 230 lines · python Raw
1 #!/usr/bin/env python3
2 """
3 Download ToolBench + APIGen-MT + ToolACE
4 and convert them to Qwen 2.5 SFT JSONL format
5 (with tool calling / function calling support).
6 """
7
8 import os
9 import json
10 import gzip
11 import tarfile
12 import zipfile
13 import requests
14 from pathlib import Path
15 from tqdm import tqdm
16 from datasets import load_dataset
17 from huggingface_hub import hf_hub_download, snapshot_download
18
19 # ====================== CONFIG ======================
20 OUTPUT_DIR = Path("./qwen25_tool_sft")
21 OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
22
23 FINAL_JSONL = OUTPUT_DIR / "tool_sft_qwen25.jsonl"
24
25 # ====================================================
26
27 def download_file(url: str, dest: Path):
28 if dest.exists():
29 print(f"[skip] {dest.name} already exists")
30 return
31 print(f"Downloading {url} ...")
32 with requests.get(url, stream=True) as r:
33 r.raise_for_status()
34 total = int(r.headers.get("content-length", 0))
35 with open(dest, "wb") as f, tqdm(total=total, unit="B", unit_scale=True) as pbar:
36 for chunk in r.iter_content(chunk_size=8192):
37 f.write(chunk)
38 pbar.update(len(chunk))
39
40
41 def to_qwen_messages(system: str | None, conversations: list[dict]) -> dict:
42 """
43 Convert a list of turns into Qwen 2.5 messages format.
44 conversations: list of {"from": "human/gpt/function/...", "value": "..."}
45 """
46 messages = []
47 if system:
48 messages.append({"role": "system", "content": system})
49
50 for turn in conversations:
51 role = turn.get("from", "").lower()
52 content = turn.get("value", "").strip()
53 if not content:
54 continue
55
56 if role in ("human", "user"):
57 messages.append({"role": "user", "content": content})
58 elif role in ("gpt", "assistant"):
59 messages.append({"role": "assistant", "content": content})
60 elif role in ("function", "tool", "observation"):
61 # Qwen-style tool response
62 messages.append({"role": "tool", "content": content})
63 else:
64 # fallback
65 messages.append({"role": "user", "content": content})
66
67 return {"messages": messages}
68
69
70 # ----------------------------------------------------
71 # 1. ToolBench (official)
72 # ----------------------------------------------------
73 def process_toolbench():
74 print("\n=== ToolBench ===")
75 # ToolBench is available on Hugging Face
76 try:
77 ds = load_dataset("ToolBench/ToolBench", split="train", trust_remote_code=True)
78 except Exception:
79 # fallback to the processed version that many people use
80 ds = load_dataset("lmsys/toolbench", split="train")
81
82 count = 0
83 with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
84 for sample in tqdm(ds, desc="ToolBench"):
85 # ToolBench usually has "conversations" or "messages"
86 convs = sample.get("conversations") or sample.get("messages") or []
87 if not convs:
88 continue
89
90 # Some versions already have role/content
91 if isinstance(convs[0], dict) and "role" in convs[0]:
92 messages = []
93 for m in convs:
94 role = m.get("role", "user")
95 content = m.get("content", "")
96 if role == "function":
97 role = "tool"
98 messages.append({"role": role, "content": content})
99 record = {"messages": messages}
100 else:
101 record = to_qwen_messages(None, convs)
102
103 if len(record["messages"]) >= 2:
104 fout.write(json.dumps(record, ensure_ascii=False) + "\n")
105 count += 1
106 print(f"ToolBench → {count} samples")
107
108
109 # ----------------------------------------------------
110 # 2. APIGen-MT (multi-turn tool calling)
111 # ----------------------------------------------------
112 def process_apigen_mt():
113 print("\n=== APIGen-MT ===")
114 # Common locations / names
115 possible = [
116 "Salesforce/APIGen-MT",
117 "Salesforce/xLAM-APIGen",
118 "Salesforce/APIGen",
119 ]
120 ds = None
121 for name in possible:
122 try:
123 ds = load_dataset(name, split="train")
124 print(f"Loaded {name}")
125 break
126 except Exception:
127 continue
128
129 if ds is None:
130 print("APIGen-MT not found on HF under common names. Skipping.")
131 return
132
133 count = 0
134 with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
135 for sample in tqdm(ds, desc="APIGen-MT"):
136 # APIGen usually has "messages" already close to OpenAI format
137 messages = sample.get("messages") or sample.get("conversations")
138 if not messages:
139 continue
140
141 # Normalize role names
142 normalized = []
143 for m in messages:
144 role = m.get("role", "user").lower()
145 content = m.get("content", "")
146 if role == "function":
147 role = "tool"
148 normalized.append({"role": role, "content": content})
149
150 if len(normalized) >= 2:
151 fout.write(json.dumps({"messages": normalized}, ensure_ascii=False) + "\n")
152 count += 1
153 print(f"APIGen-MT → {count} samples")
154
155
156 # ----------------------------------------------------
157 # 3. ToolACE
158 # ----------------------------------------------------
159 def process_toolace():
160 print("\n=== ToolACE ===")
161 possible = [
162 "Team-ACE/ToolACE",
163 "ToolACE/ToolACE",
164 "microsoft/ToolACE",
165 ]
166 ds = None
167 for name in possible:
168 try:
169 ds = load_dataset(name, split="train")
170 print(f"Loaded {name}")
171 break
172 except Exception:
173 continue
174
175 if ds is None:
176 print("ToolACE not found under common names. Trying alternative...")
177 # Some people host processed versions
178 try:
179 ds = load_dataset("json", data_files="https://huggingface.co/datasets/Team-ACE/ToolACE/resolve/main/data/train.json")
180 except Exception:
181 print("Could not load ToolACE. Skipping.")
182 return
183
184 count = 0
185 with open(FINAL_JSONL, "a", encoding="utf-8") as fout:
186 for sample in tqdm(ds, desc="ToolACE"):
187 messages = sample.get("messages") or sample.get("conversations") or []
188 if not messages:
189 continue
190
191 normalized = []
192 for m in messages:
193 if isinstance(m, dict):
194 role = m.get("role", m.get("from", "user")).lower()
195 content = m.get("content", m.get("value", ""))
196 else:
197 continue
198 if role in ("function", "observation"):
199 role = "tool"
200 elif role in ("human", "user"):
201 role = "user"
202 elif role in ("gpt", "assistant"):
203 role = "assistant"
204 normalized.append({"role": role, "content": content})
205
206 if len(normalized) >= 2:
207 fout.write(json.dumps({"messages": normalized}, ensure_ascii=False) + "\n")
208 count += 1
209 print(f"ToolACE → {count} samples")
210
211
212 # ----------------------------------------------------
213 # Main
214 # ----------------------------------------------------
215 if __name__ == "__main__":
216 # Clear previous output if you want a fresh file
217 if FINAL_JSONL.exists():
218 print(f"Removing old {FINAL_JSONL}")
219 FINAL_JSONL.unlink()
220
221 process_toolbench()
222 process_apigen_mt()
223 process_toolace()
224
225 # Final stats
226 total = sum(1 for _ in open(FINAL_JSONL, "r", encoding="utf-8"))
227 print(f"\n✅ Done! Total samples written → {FINAL_JSONL}")
228 print(f" Total lines: {total}")
229 print("\nYou can now use this JSONL for Qwen2.5 SFT (tool calling / function calling).")
230