src/render.py

# 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"<tool_call>\s*(.*?)\s*</tool_call>", re.S)
TR_RE=re.compile(r"<tool_response>\s*(.*?)\s*</tool_response>", re.S)
TOOLS_RE=re.compile(r"<tools>\s*(.*?)\s*</tools>", 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 <model_dir> [set ...], all sets when none is given
    init(sys.argv[1])
    for name in sys.argv[2:] or SETS: run(name,*SETS[name])