# render datasets into Qwen3.5 chat-template text + assistant loss spans
import json, re, os, sys, random
from multiprocessing import Pool
import jinja2
from datasets import load_from_disk
OUT="data/rendered"; os.makedirs(OUT, exist_ok=True)
def tojson(v, **kw): return json.dumps(v, ensure_ascii=False)
def to_string(v):
if isinstance(v,bool): return "true" if v else "false"
if v is None: return "null"
return str(v)
env=jinja2.Environment(extensions=["jinja2.ext.loopcontrols"], trim_blocks=False, lstrip_blocks=False)
env.filters["tojson"]=tojson; env.filters["string"]=to_string
env.globals["raise_exception"]=lambda m: (_ for _ in ()).throw(Exception(m))
TPL=None
def init(model_dir):
# the chat template ships with the model, so it comes from the same directory as the
# tokenizer and the image processor
global TPL
TPL=env.from_string(open(os.path.join(model_dir,"chat_template.jinja")).read())
SPAN_RE=re.compile(r"<\|im_start\|>assistant\n(.*?<\|im_end\|>)", re.S)
def render(messages, tools=None):
text=TPL.render(messages=messages, tools=tools, add_generation_prompt=False)
spans=[[m.start(1), m.end(1)] for m in SPAN_RE.finditer(text)]
return text, spans
TYPE_MAP={"str":"string","string":"string","int":"integer","integer":"integer","float":"number","number":"number","bool":"boolean","boolean":"boolean","list":"array","array":"array","dict":"object","object":"object"}
def xlam_tool_to_oai(t):
if t.get("type")=="function" and isinstance(t.get("function"),dict): return t
props={}; req=[]
params=t.get("parameters") or {}
if isinstance(params,dict) and "properties" in params: return {"type":"function","function":{"name":t["name"],"description":t.get("description",""),"parameters":params}}
for name,spec in (params.items() if isinstance(params,dict) else []):
spec=spec if isinstance(spec,dict) else {}
ty=str(spec.get("type","str")).replace(", optional","").strip().lower()
base=ty.split("[")[0]
js={"type":TYPE_MAP.get(base,"string")}
if base in("list","array"): js["items"]={"type":TYPE_MAP.get(ty[ty.find("[")+1:ty.rfind("]")].lower(),"string")} if "[" in ty else {"type":"string"}
if spec.get("description"): js["description"]=spec["description"]
props[name]=js
if "default" not in spec: req.append(name)
return {"type":"function","function":{"name":t["name"],"description":t.get("description",""),"parameters":{"type":"object","properties":props,"required":req}}}
def norm_apigen(r):
try: tools=json.loads(r["tools"]); ans=json.loads(r["answers"])
except Exception: return None
if not isinstance(ans,list) or not ans or not isinstance(tools,list) or not tools: return None
for a in ans:
if not isinstance(a,dict) or "name" not in a or not isinstance(a.get("arguments"),dict): return None
tools=[xlam_tool_to_oai(t) for t in tools]
msgs=[{"role":"user","content":r["query"]},{"role":"assistant","content":"","tool_calls":[{"type":"function","function":{"name":a["name"],"arguments":a["arguments"]}} for a in ans]}]
return msgs,tools
TC_RE=re.compile(r"\s*(.*?)\s*", re.S)
TR_RE=re.compile(r"\s*(.*?)\s*", re.S)
TOOLS_RE=re.compile(r"\s*(.*?)\s*", re.S)
def norm_hermes(r):
conv=r["conversations"]; tools=None
if r.get("tools"):
try: tools=json.loads(r["tools"])
except Exception: tools=None
msgs=[]
for m in conv:
f,v=m["from"],m["value"]
if f=="system":
if tools is None:
mm=TOOLS_RE.search(v)
if mm:
try: tools=json.loads(mm.group(1))
except Exception: return None
continue
if f=="human": msgs.append({"role":"user","content":v})
elif f=="gpt":
calls=[]
for c in TC_RE.findall(v):
try: d=json.loads(c)
except Exception: return None
if "name" not in d: return None
calls.append({"type":"function","function":{"name":d["name"],"arguments":d.get("arguments",{})}})
content=TC_RE.sub("",v).strip()
msg={"role":"assistant","content":content}
if calls: msg["tool_calls"]=calls
msgs.append(msg)
elif f=="tool":
resps=TR_RE.findall(v) or [v.strip()]
for x in resps: msgs.append({"role":"tool","content":x})
else: return None
if tools is not None:
if not isinstance(tools,list): return None
tools=[t if t.get("type")=="function" else {"type":"function","function":t} for t in tools]
if not msgs or msgs[0]["role"]!="user": return None
return msgs,tools
def norm_nemotron(r):
# OpenAI style trajectories with reasoning_content dropped and tool outputs serialized as JSON text
msgs=[]
for m in r["messages"]:
role=m["role"]; c=m.get("content")
if role=="system":
if c: msgs.append({"role":"system","content":c})
elif role=="user": msgs.append({"role":"user","content":c})
elif role=="assistant":
msg={"role":"assistant","content":c or ""}
if m.get("tool_calls"):
calls=[]
for t in m["tool_calls"]:
a=t["function"]["arguments"]
if isinstance(a,str): a=json.loads(a)
if not isinstance(a,dict): return None
calls.append({"type":"function","function":{"name":t["function"]["name"],"arguments":a}})
msg["tool_calls"]=calls
msgs.append(msg)
elif role=="tool": msgs.append({"role":"tool","content":c if isinstance(c,str) else json.dumps(c)})
else: return None
if not msgs or msgs[-1]["role"]!="assistant" or msgs[-1].get("tool_calls"): return None
n_calls=sum(1 for m in msgs if m["role"]=="assistant" and m.get("tool_calls"))
if n_calls==0 or (n_calls==1 and random.random()>0.15): return None
return msgs,r["tools"]
def norm_messages(r):
msgs=r["messages"]
if not msgs or msgs[0]["role"] not in("user","system"): return None
return msgs,None
def work(args):
fn,r=args
try:
x=fn(r)
if x is None: return None
text,spans=render(*x)
except Exception as e:
return None
if not spans: return None
return json.dumps({"text":text,"spans":spans},ensure_ascii=False)
def run(name,path,split,fn,holdout=0.0):
if path.endswith(".jsonl"): ds=[json.loads(l) for l in open(path)]
else:
ds=load_from_disk(path); ds=ds[split] if split in ds else ds["train"]
random.seed(1)
rows=[(fn,ds[i]) for i in range(len(ds))]
with Pool(18) as p: out=[x for x in p.imap(work,rows,chunksize=256) if x]
random.seed(0); random.shuffle(out)
n_ev=int(len(out)*holdout)
with open(f"{OUT}/{name}.jsonl","w") as f: f.write("\n".join(out[n_ev:])+"\n")
if n_ev:
with open(f"{OUT}/{name}.eval.jsonl","w") as f: f.write("\n".join(out[:n_ev])+"\n")
print(f"{name}: {len(ds)} -> {len(out)} kept ({n_ev} eval)",flush=True)
SETS={
"smoltalk":("data/smol_smoltalk","train",norm_messages,0.0),
"smoltalk_eval":("data/smol_smoltalk","test",norm_messages,0.0),
"everyday":("data/everyday","train",norm_messages,0.05),
"apigen":("data/apigen","train",norm_apigen,0.01),
"hermes_fc":("data/hermes_fc","train",norm_hermes,0.05),
"hermes_fc_single":("data/hermes_fc_single","train",norm_hermes,0.05),
"hermes_glaive":("data/hermes_glaive","train",norm_hermes,0.05),
"nemotron":("data/nemotron/data/tool_calling.jsonl","train",norm_nemotron,0.005),
}
if __name__=="__main__":
# usage: render.py [set ...], all sets when none is given
init(sys.argv[1])
for name in sys.argv[2:] or SETS: run(name,*SETS[name])