feat(training): add the colab finetuning pipeline

Reproduces these weights end to end, and carries the chat template the
adapter cannot work without.

train_sahabatai.py imports its data helpers from train_qwen.py because
that is what ran on the VM; the sibling adapter's repo ships the same
file. check_template.py proves the template off-GPU -- it fails by
producing plausible text rather than raising, so catching it after a
five-hour run is the expensive way to find out.

training/check_template.py

@@ -0,0 +1,138 @@
+"""Validate the Llama-3 tool template against the real corpus, off-GPU.
+
+Everything here runs on jinja2 + tokenizers alone -- no torch, no VM -- so the
+template is proven before any GPU time is spent. It checks the two things that
+fail silently rather than loudly:
+
+  * tool calls and tool schemas survive rendering (the stock Sahabat-AI
+    template drops both, and an assistant tool-call turn has `content: ""`, so
+    the damage looks like an ordinary empty reply);
+  * response-only masking covers exactly the assistant turns. A tool result
+    rendered under a header the instruction delimiter does not match would sit
+    inside the loss and teach the model to write its own forecasts.
+
+    python3 training/check_template.py --data train.parquet \
+        --tokenizer tokenizer.json
+"""
+
+import argparse
+import json
+import statistics
+from pathlib import Path
+
+import jinja2.ext
+import pyarrow.parquet as pq
+from jinja2.sandbox import ImmutableSandboxedEnvironment
+
+TEMPLATE = Path(__file__).resolve().parent / "llama3_tools.jinja"
+
+BOS = "<|begin_of_text|>"
+INSTRUCTION_PART = "<|start_header_id|>user<|end_header_id|>\n\n"
+RESPONSE_PART = "<|start_header_id|>assistant<|end_header_id|>\n\n"
+
+
+def build_env():
+    """Reproduce transformers' template environment closely enough to trust."""
+    env = ImmutableSandboxedEnvironment(
+        trim_blocks=True, lstrip_blocks=True, extensions=[jinja2.ext.loopcontrols])
+    env.filters["tojson"] = lambda x, **kw: json.dumps(
+        x, ensure_ascii=kw.get("ensure_ascii", False), indent=kw.get("indent"))
+    env.globals["raise_exception"] = lambda m: (_ for _ in ()).throw(Exception(m))
+    return env
+
+
+def hydrate(row):
+    messages = []
+    for message in row["messages"]:
+        clean = {k: v for k, v in message.items() if v is not None}
+        if "tool_calls" in clean:
+            clean["tool_calls"] = json.loads(clean["tool_calls"])
+        messages.append(clean)
+    tools = row.get("tools")
+    return messages, (json.loads(tools) if tools else None)
+
+
+def trained_spans(text):
+    """The character spans train_on_responses_only would keep in the loss.
+
+    unsloth trains from each response delimiter to the next *instruction*
+    delimiter, or to the end. Working in characters rather than token ids is
+    equivalent here and keeps the check tokenizer-free.
+    """
+    spans, cursor = [], 0
+    while (start := text.find(RESPONSE_PART, cursor)) != -1:
+        start += len(RESPONSE_PART)
+        stop = text.find(INSTRUCTION_PART, start)
+        stop = len(text) if stop == -1 else stop
+        spans.append((start, stop))
+        cursor = stop
+    return spans
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data", required=True,
+                        help="train.parquet from prepare_data.py")
+    parser.add_argument("--tokenizer", help="path to a tokenizer.json, for length stats")
+    parser.add_argument("--rows", type=int, default=4000)
+    args = parser.parse_args()
+
+    template = build_env().from_string(TEMPLATE.read_text(encoding="utf-8"))
+    rows = pq.read_table(args.data).slice(0, args.rows).to_pylist()
+
+    calls = schemas = leaked = empty_mask = 0
+    tool_rows = 0
+    lengths = []
+    tokenizer = None
+    if args.tokenizer:
+        from tokenizers import Tokenizer
+        tokenizer = Tokenizer.from_file(args.tokenizer)
+
+    for row in rows:
+        messages, tools = hydrate(row)
+        text = template.render(messages=messages, tools=tools, bos_token=BOS,
+                               add_generation_prompt=False)
+
+        spans = trained_spans(text)
+        if not spans:
+            empty_mask += 1
+        trained = "".join(text[a:b] for a, b in spans)
+
+        # A tool result inside the trained span is the failure this exists for.
+        leaked += trained.count("<tool_response>")
+
+        for message in messages:
+            for call in message.get("tool_calls") or []:
+                tool_rows += 1
+                needle = f'{{"name": "{call["function"]["name"]}", "parameters": '
+                calls += needle in text
+                # ...and the call itself MUST be trained on, or the model never
+                # learns to emit one.
+                calls -= needle not in trained
+        if tools:
+            schemas += all(t["function"]["description"][:40] in text for t in tools)
+
+        if tokenizer:
+            lengths.append(len(tokenizer.encode(text, add_special_tokens=False).ids))
+
+    print(f"rows                 {len(rows)}")
+    print(f"tool calls           {calls}/{tool_rows} rendered and inside the loss")
+    print(f"tool schemas         {schemas} conversations carry the offered schema")
+    print(f"tool results leaked  {leaked}  (must be 0)")
+    print(f"empty label masks    {empty_mask}  (must be 0)")
+
+    if lengths:
+        lengths.sort()
+        n = len(lengths)
+        print(f"tokens               mean={statistics.mean(lengths):.0f} "
+              f"p50={lengths[n // 2]} p90={lengths[int(n * .9)]} "
+              f"p99={lengths[int(n * .99)]} max={lengths[-1]} "
+              f"over_4096={sum(1 for x in lengths if x > 4096)}")
+
+    if leaked or empty_mask or calls != tool_rows:
+        raise SystemExit("template check FAILED")
+    print("template check OK")
+
+
+if __name__ == "__main__":
+    main()

training/drive.py

@@ -0,0 +1,373 @@
+"""Drive a Colab VM over the REST API, independently of the CLI session record.
+
+    ~/.local/share/uv/tools/google-colab-cli/bin/python colab/drive.py <cmd> ...
+
+Run it with the colab-cli tool env's own interpreter; a system python3 fails on
+pydantic_core.
+
+Why this exists: the runtime proxy token has a 1-hour TTL. When it expires the
+CLI prunes the local session record and kills the keep-alive daemon with it, and
+the VM is reclaimed roughly 40 minutes later -- long after the event that doomed
+it. `colab upload` then reports `File or directory not found` for a remote path
+that is perfectly fine; that error is the symptom of the prune, not evidence
+about the VM. So once a job is launched, nothing here routes through
+`colab -s <name>`: every call re-mints credentials, calls keep_alive before it
+does anything else, and depends on no local state.
+
+Do NOT run `colab new` to recover an orphaned session -- it provisions a second
+billable GPU rather than reattaching to the first.
+
+    keepalive                      touch the VM and report what is assigned
+    ls    <remote_dir>             list a directory
+    cat   <remote_file> [--tail N] read a text file
+    get   <remote_file> <local>    download (handles base64 payloads)
+    put   <local> <remote>         upload
+    run   <local_py> [--timeout N] execute a local script on the VM
+    watch <remote_log> --done M    poll a log, refreshing keep-alive each pass
+"""
+
+import argparse
+import base64
+import os
+import shutil
+import sys
+import time
+
+import requests
+
+CONFIG = os.path.expanduser("~/.colab-cli-oauth-config.json")
+MATCH = os.environ.get("COLAB_MATCH", "a100")
+
+
+def connect():
+    """Mint fresh credentials, keep the VM alive, return its proxy info."""
+    from colab_cli.auth import AuthProvider, get_credentials
+    from colab_cli.client import Client, Prod
+
+    try:
+        creds = get_credentials(CONFIG, provider=AuthProvider.ADC)
+    except Exception:
+        creds = get_credentials(CONFIG, provider=AuthProvider.OAUTH2)
+
+    client = Client(Prod(), creds)
+    assignment = next(
+        (a for a in client.list_assignments() if MATCH.lower() in a.endpoint.lower()),
+        None,
+    )
+    if assignment is None:
+        raise SystemExit(f"no assignment matching {MATCH!r} -- the VM may be gone")
+
+    # Before anything else. Reading a log does not reset the idle timer.
+    client.keep_alive_assignment(assignment.endpoint)
+    return assignment, assignment.runtime_proxy_info
+
+
+def contents(proxy, path, method="GET", payload=None, params=None):
+    query = {"authuser": "0", "colab-runtime-proxy-token": proxy.token}
+    query.update(params or {})
+    url = f"{proxy.url}/api/contents/{path.lstrip('/')}"
+    response = requests.request(method, url, params=query, json=payload, timeout=300)
+    response.raise_for_status()
+    return response.json() if response.content else {}
+
+
+def read_text(proxy, path):
+    """Return a remote file as text, whatever format the API chose."""
+    payload = contents(proxy, path, params={"content": "1"})
+    if payload.get("format") == "base64":
+        return base64.b64decode(payload["content"]).decode("utf-8", "replace")
+    return payload.get("content", "")
+
+
+def cmd_keepalive(args, assignment, proxy):
+    print(f"assignment {assignment.endpoint}")
+    print(f"proxy      {proxy.url}")
+    print("keep-alive ok")
+
+
+def cmd_ls(args, assignment, proxy):
+    listing = contents(proxy, args.path, params={"content": "1"})
+    for entry in sorted(listing.get("content", []), key=lambda c: c["name"]):
+        size = entry.get("size")
+        print(f"  {entry['type']:<9} {str(size if size is not None else '-'):>12}  "
+              f"{entry['name']}")
+
+
+def cmd_cat(args, assignment, proxy):
+    text = read_text(proxy, args.path)
+    lines = text.splitlines()
+    for line in (lines[-args.tail:] if args.tail else lines):
+        print(line)
+
+
+def download(proxy, remote, local):
+    """Stream a remote file to disk.
+
+    Jupyter's /files/ endpoint serves raw bytes. The contents API would wrap a
+    170 MB checkpoint in base64 inside a JSON document -- a third larger on the
+    wire and fully buffered in memory at both ends -- so that is only the
+    fallback.
+    """
+    os.makedirs(os.path.dirname(os.path.abspath(local)) or ".", exist_ok=True)
+    query = {"authuser": "0", "colab-runtime-proxy-token": proxy.token}
+    url = f"{proxy.url}/files/{remote.lstrip('/')}"
+
+    try:
+        with requests.get(url, params=query, stream=True, timeout=900) as response:
+            response.raise_for_status()
+            written = 0
+            with open(local, "wb") as handle:
+                for chunk in response.iter_content(1 << 20):
+                    handle.write(chunk)
+                    written += len(chunk)
+        return written
+    except requests.HTTPError:
+        payload = contents(proxy, remote, params={"content": "1"})
+        blob = (base64.b64decode(payload["content"])
+                if payload.get("format") == "base64"
+                else payload.get("content", "").encode())
+        with open(local, "wb") as handle:
+            handle.write(blob)
+        return len(blob)
+
+
+def cmd_get(args, assignment, proxy):
+    written = download(proxy, args.remote, args.local)
+    print(f"{args.remote} -> {args.local} ({written} bytes)")
+
+
+def newest_checkpoint(proxy, remote_dir):
+    listing = contents(proxy, remote_dir, params={"content": "1"})
+    checkpoints = [
+        (int(c["name"].split("-")[1]), c["name"])
+        for c in listing.get("content", [])
+        if c["type"] == "directory" and c["name"].startswith("checkpoint-")
+        and c["name"].split("-")[-1].isdigit()
+    ]
+    return max(checkpoints) if checkpoints else (None, None)
+
+
+def cmd_pull(args, assignment, proxy):
+    """Copy the newest checkpoint down, so a VM reclaim costs minutes not hours.
+
+    Checkpoints live only on the VM. Keep-alive makes a reclaim unlikely, not
+    impossible, and save_total_limit prunes the VM's own copies as it goes.
+    """
+    deadline = time.time() + args.max_seconds if args.max_seconds else None
+    os.makedirs(args.out_dir, exist_ok=True)
+    pulled = set()
+
+    while True:
+        try:
+            if pulled:
+                assignment, proxy = connect()
+
+            step, name = newest_checkpoint(proxy, args.remote_dir)
+            log = read_text(proxy, args.log) if args.log else ""
+
+            if name and name not in pulled:
+                target = os.path.join(args.out_dir, name)
+                listing = contents(proxy, f"{args.remote_dir}/{name}",
+                                   params={"content": "1"})
+                total = 0
+                for entry in listing.get("content", []):
+                    if entry["type"] != "file":
+                        continue
+                    got = download(proxy, f"{args.remote_dir}/{name}/{entry['name']}",
+                                   os.path.join(target, entry["name"]))
+                    if entry.get("size") and got != entry["size"]:
+                        raise IOError(f"{entry['name']}: got {got} of {entry['size']}")
+                    total += got
+                pulled.add(name)
+                print(f"[pull] {time.strftime('%H:%M:%S')} {name} "
+                      f"-> {target} ({total/1e6:.0f} MB)", flush=True)
+
+                # Keep only the newest few; each is ~277 MB.
+                local = sorted(
+                    (int(d.split("-")[1]), d) for d in os.listdir(args.out_dir)
+                    if d.startswith("checkpoint-") and d.split("-")[-1].isdigit())
+                for _, old in local[:-args.keep]:
+                    shutil.rmtree(os.path.join(args.out_dir, old), ignore_errors=True)
+                    print(f"[pull] pruned local {old}", flush=True)
+            else:
+                print(f"[pull] {time.strftime('%H:%M:%S')} newest={name} "
+                      f"already held, keep-alive ok", flush=True)
+
+            if args.done and args.done in log:
+                print(f"[pull] saw {args.done}", flush=True)
+                return 0
+        except Exception as exc:
+            print(f"[pull] retry: {type(exc).__name__}: {str(exc)[:140]}", flush=True)
+
+        if deadline and time.time() > deadline:
+            return 0
+        time.sleep(args.interval)
+
+
+# A single base64 PUT of a large file is accepted and then silently dropped:
+# 170 MB of weights became a 227 MB JSON body and never landed, while the call
+# returned success. Anything above this goes up in parts and is verified.
+PUT_CHUNK = 24 * 1024 * 1024
+
+
+def remote_size(proxy, path):
+    try:
+        return contents(proxy, path, params={"content": "0"}).get("size")
+    except requests.HTTPError:
+        return None
+
+
+def put_blob(proxy, blob, remote):
+    contents(proxy, remote, method="PUT", payload={
+        "type": "file", "format": "base64",
+        "content": base64.b64encode(blob).decode(),
+    })
+
+
+def cmd_put(args, assignment, proxy):
+    blob = open(args.local, "rb").read()
+
+    if len(blob) <= PUT_CHUNK:
+        put_blob(proxy, blob, args.remote)
+        got = remote_size(proxy, args.remote)
+        if got != len(blob):
+            raise SystemExit(f"upload verify failed: remote {got} of {len(blob)} bytes")
+        print(f"{args.local} -> {args.remote} ({len(blob)} bytes, verified)")
+        return
+
+    parts = [blob[i:i + PUT_CHUNK] for i in range(0, len(blob), PUT_CHUNK)]
+    for index, part in enumerate(parts):
+        name = f"{args.remote}.part{index:03d}"
+        assignment, proxy = connect()          # each part re-mints and keeps alive
+        put_blob(proxy, part, name)
+        got = remote_size(proxy, name)
+        if got != len(part):
+            raise SystemExit(f"part {index} failed: remote {got} of {len(part)} bytes")
+        print(f"  part {index + 1}/{len(parts)} ok ({len(part)} bytes)", flush=True)
+
+    print(f"{args.local} -> {args.remote} ({len(blob)} bytes in {len(parts)} parts)")
+    print("  now run:  drive.py join <remote> <nparts> <expected_bytes>")
+
+
+def cmd_join(args, assignment, proxy):
+    """Concatenate uploaded parts on the VM and verify the result."""
+    from colab_cli.runtime import ColabRuntime
+
+    code = (
+        "import os\n"
+        f"target = {args.remote!r}\n"
+        f"n = {args.nparts}\n"
+        "with open(target, 'wb') as out:\n"
+        "    for i in range(n):\n"
+        "        p = f'{target}.part{i:03d}'\n"
+        "        with open(p, 'rb') as fh:\n"
+        "            out.write(fh.read())\n"
+        "        os.remove(p)\n"
+    )
+    runtime = ColabRuntime(url=proxy.url, token=proxy.token)
+    try:
+        runtime.execute_code(code, timeout=args.timeout)
+    except Exception as exc:
+        print(f"[join] execute returned {type(exc).__name__}; verifying anyway", flush=True)
+
+    deadline = time.time() + 300
+    while time.time() < deadline:
+        assignment, proxy = connect()
+        got = remote_size(proxy, args.remote)
+        if got == args.expect:
+            print(f"{args.remote} joined and verified ({got} bytes)")
+            return 0
+        time.sleep(10)
+
+    raise SystemExit(f"join failed: remote {remote_size(proxy, args.remote)} "
+                     f"of {args.expect} bytes")
+
+
+def cmd_run(args, assignment, proxy):
+    from colab_cli.runtime import ColabRuntime
+
+    runtime = ColabRuntime(url=proxy.url, token=proxy.token)
+    code = open(args.script, encoding="utf-8").read()
+    for reply in runtime.execute_code(code, timeout=args.timeout):
+        for key in ("text", "traceback", "evalue"):
+            if key in reply:
+                value = reply[key]
+                print("".join(value) if isinstance(value, list) else value, end="")
+    print()
+
+
+def cmd_watch(args, assignment, proxy):
+    """Poll a log. Every pass reconnects, so no single expiry ends the watch."""
+    deadline = time.time() + args.max_seconds if args.max_seconds else None
+    seen = 0
+
+    while True:
+        try:
+            if seen:                                    # first pass already connected
+                assignment, proxy = connect()
+            text = read_text(proxy, args.path)
+        except Exception as exc:
+            # An unreadable poll is a retry, not a failure. The job outlives it.
+            print(f"[watch] retry: {type(exc).__name__}: {str(exc)[:120]}", flush=True)
+            time.sleep(args.interval)
+            continue
+
+        seen += 1
+        lines = text.splitlines()
+        print(f"--- {time.strftime('%H:%M:%S')}  {len(lines)} lines  keep-alive ok",
+              flush=True)
+        for line in lines[-args.tail:]:
+            print(line, flush=True)
+
+        for marker in filter(None, [args.done, args.fail]):
+            if marker in text:
+                print(f"[watch] saw {marker}", flush=True)
+                return 0 if marker == args.done else 1
+
+        if deadline and time.time() > deadline:
+            print("[watch] time budget reached; the job is unaffected", flush=True)
+            return 0
+        time.sleep(args.interval)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    sub.add_parser("keepalive")
+
+    p = sub.add_parser("ls"); p.add_argument("path", nargs="?", default="/content")
+    p = sub.add_parser("cat"); p.add_argument("path"); p.add_argument("--tail", type=int, default=0)
+    p = sub.add_parser("get"); p.add_argument("remote"); p.add_argument("local")
+    p = sub.add_parser("put"); p.add_argument("local"); p.add_argument("remote")
+    p = sub.add_parser("run"); p.add_argument("script"); p.add_argument("--timeout", type=float, default=300)
+
+    p = sub.add_parser("join")
+    p.add_argument("remote"); p.add_argument("nparts", type=int)
+    p.add_argument("expect", type=int); p.add_argument("--timeout", type=float, default=120)
+
+    p = sub.add_parser("pull")
+    p.add_argument("--remote-dir", default="/content/outputs/qwen3_5_4b_fairleap")
+    p.add_argument("--out-dir", default="checkpoints")
+    p.add_argument("--log", default="/content/logs/train.log")
+    p.add_argument("--interval", type=int, default=300)
+    p.add_argument("--keep", type=int, default=2)
+    p.add_argument("--done", default="TRAINING_COMPLETE")
+    p.add_argument("--max-seconds", type=int, default=0)
+
+    p = sub.add_parser("watch")
+    p.add_argument("path")
+    p.add_argument("--tail", type=int, default=20)
+    p.add_argument("--interval", type=int, default=60)
+    p.add_argument("--done", default="TRAINING_COMPLETE")
+    p.add_argument("--fail", default="")
+    p.add_argument("--max-seconds", type=int, default=0)
+
+    args = parser.parse_args()
+    assignment, proxy = connect()
+    handler = globals()[f"cmd_{args.cmd}"]
+    return handler(args, assignment, proxy) or 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

training/llama3_tools.jinja

@@ -0,0 +1,63 @@
+{#- Tool-capable Llama-3 chat template for the Fairleap corpus.
+
+    The stock Sahabat-AI template renders every message as
+    `content | trim` and ignores both `tool_calls` and the `tools` argument.
+    On this corpus that is silent data loss, not an error: an assistant
+    tool-call turn carries `content: ""`, so it renders as an *empty*
+    assistant reply and the call disappears. ~10% of the corpus is shaped
+    that way.
+
+    Two things this adds:
+      1. assistant `tool_calls` render as the one-line JSON the model is
+         meant to emit, and `role: tool` results come back as a `user` turn
+         wrapping `<tool_response>` so response-only masking still works
+         (see the note on that branch);
+      2. the `tools` schema list is folded into the system turn, so
+         "tool offered" and "tool not offered" are distinguishable in
+         context. Without that the model cannot learn when *not* to call.
+-#}
+{{- bos_token }}
+{%- set has_system = messages and messages[0]['role'] == 'system' %}
+{%- set tool_header = 'Kamu punya akses ke fungsi berikut. Untuk memanggil fungsi, balas HANYA dengan satu baris JSON berbentuk {"name": <nama fungsi>, "parameters": <objek argumen>}, tanpa teks lain. Panggil fungsi hanya jika pertanyaan driver memang membutuhkannya.\n\nFungsi yang tersedia:' %}
+{%- if tools and not has_system %}
+    {{- '<|start_header_id|>system<|end_header_id|>\n\n' + tool_header }}
+    {%- for tool in tools %}
+        {{- '\n' }}{{- tool['function'] | tojson }}
+    {%- endfor %}
+    {{- '<|eot_id|>' }}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message['role'] == 'system' %}
+        {{- '<|start_header_id|>system<|end_header_id|>\n\n' + message['content'] | trim }}
+        {%- if tools and loop.first %}
+            {{- '\n\n' + tool_header }}
+            {%- for tool in tools %}
+                {{- '\n' }}{{- tool['function'] | tojson }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|eot_id|>' }}
+    {%- elif message['role'] == 'assistant' and message.get('tool_calls') %}
+        {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+        {%- for tool_call in message['tool_calls'] %}
+            {#- Emit each piece separately: `tojson` returns Markup, and
+                concatenating a plain string with it HTML-escapes the quotes. #}
+            {{- '{"name": "' }}{{- tool_call['function']['name'] }}{{- '", "parameters": ' }}
+            {{- tool_call['function']['arguments'] | tojson }}{{- '}' }}
+        {%- endfor %}
+        {{- '<|eot_id|>' }}
+    {%- elif message['role'] == 'tool' %}
+        {#- Deliberately a `user` turn wrapping `<tool_response>`, not Llama-3.1's
+            `ipython` header. `train_on_responses_only` masks from the response
+            delimiter to the next *instruction* delimiter, and an `ipython`
+            header matches neither -- the tool result would land inside the loss
+            and teach the model to invent forecasts. This is also byte-for-byte
+            the shape Qwen's own template uses, so both Fairleap adapters take
+            tool results in the same form. #}
+        {{- '<|start_header_id|>user<|end_header_id|>\n\n<tool_response>\n' + message['content'] | trim + '\n</tool_response><|eot_id|>' }}
+    {%- else %}
+        {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }}
+{%- endif %}

training/prepare_data.py

@@ -0,0 +1,131 @@
+"""Prepare the Fairleap corpus for a Colab finetuning run.
+
+Truncates conversations at their first canned follow-up answer, then writes
+Parquet with an explicit schema. Parquet because the JSONL is 286 MiB and the
+upload crosses a metered link; the explicit schema because only ~10% of records
+carry `tools`/`tool_calls` and letting Arrow infer that from a first block is
+what makes `load_dataset("json", ...)` fail on a corpus shaped like this.
+
+    python3 colab/prepare_data.py
+
+Writes colab/data/{train,validation}.parquet. train_qwen.py reads them back.
+"""
+
+import json
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+ROOT = Path(__file__).resolve().parent.parent
+SPLIT_DIR = ROOT / "data" / "splits"
+OUT_DIR = Path(__file__).resolve().parent / "data"
+
+SPLITS = ("train", "validation")
+
+# 79.8% of offline follow-up turns are this one sentence, varying only by
+# honorific, and it acknowledges rather than answers. Training on it teaches a
+# model to do the same. See the dataset card, Limitations.
+CANNED_FOLLOWUP = "Kalau ada bagian yang mau saya rinci lagi"
+
+MESSAGE_TYPE = pa.struct([
+    ("role", pa.string()),
+    ("content", pa.string()),
+    ("tool_calls", pa.string()),
+    ("name", pa.string()),
+    ("tool_call_id", pa.string()),
+])
+SCHEMA = pa.schema([
+    ("messages", pa.list_(MESSAGE_TYPE)),
+    ("tools", pa.string()),
+    ("scenario", pa.string()),
+    ("language", pa.string()),
+    ("source", pa.string()),
+])
+
+
+def truncate_canned(record, min_messages=3):
+    """Cut the conversation before its first canned follow-up.
+
+    Deleting the turn in place would leave a driver's question standing
+    unanswered, so everything from that turn onward goes and the remainder is
+    trimmed back to end on an assistant reply. Returns None if too little is
+    left to be a conversation.
+    """
+    messages = record["messages"]
+    cut = next(
+        (i for i, m in enumerate(messages)
+         if m["role"] == "assistant" and CANNED_FOLLOWUP in (m.get("content") or "")),
+        None,
+    )
+    if cut is None:
+        return record
+
+    trimmed = messages[:cut]
+    while trimmed and trimmed[-1]["role"] != "assistant":
+        trimmed.pop()
+    if len(trimmed) < min_messages:
+        return None
+
+    return {**record, "messages": trimmed}
+
+
+def flatten(record):
+    messages = []
+    for message in record["messages"]:
+        messages.append({
+            "role": message["role"],
+            "content": message.get("content"),
+            "tool_calls": (json.dumps(message["tool_calls"], ensure_ascii=False)
+                           if "tool_calls" in message else None),
+            "name": message.get("name"),
+            "tool_call_id": message.get("tool_call_id"),
+        })
+
+    tools = record.get("tools")
+    meta = record["meta"]
+    return {
+        "messages": messages,
+        "tools": None if tools is None else json.dumps(tools, ensure_ascii=False),
+        "scenario": meta["scenario"],
+        "language": meta["language"],
+        "source": meta["source"],
+    }
+
+
+def main():
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    for split in SPLITS:
+        source = SPLIT_DIR / f"fairleap_{split}.jsonl"
+        target = OUT_DIR / f"{split}.parquet"
+
+        rows, seen, truncated, dropped = [], 0, 0, 0
+        with open(source, encoding="utf-8") as handle:
+            for line in handle:
+                if not line.strip():
+                    continue
+                seen += 1
+                record = json.loads(line)
+                before = len(record["messages"])
+
+                record = truncate_canned(record)
+                if record is None:
+                    dropped += 1
+                    continue
+                if len(record["messages"]) != before:
+                    truncated += 1
+
+                rows.append(flatten(record))
+
+        pq.write_table(pa.Table.from_pylist(rows, schema=SCHEMA), target,
+                       compression="zstd")
+
+        size = target.stat().st_size
+        print(f"{split:<11} {seen:>6} read  {len(rows):>6} written  "
+              f"{truncated:>5} truncated  {dropped:>3} dropped  "
+              f"{size/1e6:>6.1f} MB")
+
+
+if __name__ == "__main__":
+    main()

training/train_qwen.py

@@ -0,0 +1,302 @@
+"""QLoRA finetune of Qwen3.5-4B on the Fairleap chat corpus. Runs on the VM.
+
+Launch detached so the job outlives the websocket, the proxy token and the
+local session record:
+
+    nohup python -u train_qwen.py > /content/logs/train.log 2>&1 &
+
+Chat templates are applied here, on the VM, into a single `text` column.
+Handing SFTTrainer the raw nested `messages` instead would make Arrow infer a
+schema from the first block, and only ~10% of records carry `tools` or
+`tool_calls` -- that mismatch is a mid-run crash, not a load-time error.
+"""
+
+import argparse
+import json
+import os
+import re
+import time
+from pathlib import Path
+
+MODEL_ID = "Qwen/Qwen3.5-4B"
+
+# ChatML delimiters. Everything outside the assistant turns is masked out of
+# the loss: without this the model also learns to reproduce the ~4,400-char
+# stuffed system prompt and the driver's own messages.
+INSTRUCTION_PART = "<|im_start|>user\n"
+RESPONSE_PART = "<|im_start|>assistant\n"
+
+SEED = 20260819
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data-dir", default="/content/data")
+    p.add_argument("--output-dir", default="/content/outputs/qwen3_5_4b_fairleap")
+    p.add_argument("--adapter-dir", default="/content/fairleap-v1-clm-qwen3.5-4b-adapter")
+    p.add_argument("--epochs", type=float, default=2.0)
+    p.add_argument("--max-seq", type=int, default=4096)
+    p.add_argument("--batch", type=int, default=2)
+    p.add_argument("--accum", type=int, default=8)
+    p.add_argument("--lr", type=float, default=2e-4)
+    p.add_argument("--save-steps", type=int, default=400)
+    p.add_argument("--eval-steps", type=int, default=200)
+    p.add_argument("--limit", type=int, default=0, help="cap records, for smoke runs")
+    p.add_argument("--sample", type=int, default=0,
+                   help="stratified subsample of train, preserving the mix")
+    p.add_argument("--group-by-length", action="store_true",
+                   help="batch similar lengths together to cut padding waste")
+    p.add_argument("--no-grad-checkpoint", action="store_true",
+                   help="trade VRAM for speed; only if the card has headroom")
+    p.add_argument("--resume", action="store_true", help="continue from newest checkpoint")
+    p.add_argument("--check", action="store_true", help="render and report, do not train")
+    return p.parse_args()
+
+
+def get_tokenizer(processor):
+    """Qwen3.5-4B is a vision-language model, so from_pretrained hands back a
+    Qwen3VLProcessor rather than a tokenizer. Its __call__ reads the first
+    positional argument as an image source, which turns every training string
+    into `Incorrect image source`. The inner text tokenizer is what the corpus
+    needs, and it carries the chat template too."""
+    tokenizer = getattr(processor, "tokenizer", processor)
+    probe = tokenizer("ping")
+    ids = probe["input_ids"] if hasattr(probe, "keys") else probe
+    if ids and isinstance(ids[0], list):
+        raise SystemExit("tokenizer returns batched ids -- still a processor, not a tokenizer")
+    return tokenizer
+
+
+def render(tokenizer, messages, tools):
+    """Render one conversation. Thinking is disabled so the corpus, which has
+    no reasoning traces, does not teach the model to open a <think> block on
+    some turns and not others."""
+    try:
+        return tokenizer.apply_chat_template(
+            messages, tools=tools, tokenize=False,
+            add_generation_prompt=False, enable_thinking=False)
+    except TypeError:
+        return tokenizer.apply_chat_template(
+            messages, tools=tools, tokenize=False, add_generation_prompt=False)
+
+
+def hydrate(row):
+    """Turn a Parquet row back into native chat-message dicts."""
+    messages = []
+    for message in row["messages"]:
+        clean = {k: v for k, v in message.items() if v is not None}
+        if "tool_calls" in clean:
+            clean["tool_calls"] = json.loads(clean["tool_calls"])
+        messages.append(clean)
+
+    tools = row.get("tools")
+    return messages, (json.loads(tools) if tools else None)
+
+
+def stratified_sample(rows, n, seed=SEED):
+    """Take n rows preserving the scenario x language x source mix.
+
+    A head slice would over-represent whatever the split happened to order
+    first; proportional allocation with largest-remainder keeps every one of
+    the 17 scenarios and both minority languages present at their real rate.
+    """
+    import random
+    from collections import defaultdict
+
+    if not n or n >= len(rows):
+        return rows
+
+    buckets = defaultdict(list)
+    for index, row in enumerate(rows):
+        buckets[(row["scenario"], row["language"], row["source"])].append(index)
+
+    exact = {k: len(v) * n / len(rows) for k, v in buckets.items()}
+    quota = {k: int(v) for k, v in exact.items()}
+    for key in sorted(exact, key=lambda k: exact[k] - quota[k], reverse=True)[
+            :n - sum(quota.values())]:
+        quota[key] += 1
+
+    rng = random.Random(seed)
+    picked = []
+    for key, indices in buckets.items():
+        picked.extend(rng.sample(indices, min(quota[key], len(indices))))
+
+    picked.sort()
+    return [rows[i] for i in picked]
+
+
+def build_dataset(tokenizer, path, limit=0, sample=0):
+    """Render every conversation to one string, so the column type is uniform."""
+    import pyarrow.parquet as pq
+    from datasets import Dataset
+
+    rows = pq.read_table(path).to_pylist()
+    if sample:
+        rows = stratified_sample(rows, sample)
+    if limit:
+        rows = rows[:limit]
+
+    texts = [render(tokenizer, *hydrate(row)) for row in rows]
+    return Dataset.from_dict({"text": texts})
+
+
+def newest_checkpoint(output_dir):
+    path = Path(output_dir)
+    if not path.is_dir():
+        return None
+    checkpoints = [
+        (int(m.group(1)), str(d))
+        for d in path.iterdir()
+        if d.is_dir() and (m := re.fullmatch(r"checkpoint-(\d+)", d.name))
+    ]
+    return max(checkpoints)[1] if checkpoints else None
+
+
+def main():
+    args = parse_args()
+
+    import torch
+    from unsloth import FastLanguageModel
+    from unsloth.chat_templates import train_on_responses_only
+    from trl import SFTConfig, SFTTrainer
+
+    bf16 = torch.cuda.is_bf16_supported()
+    name = torch.cuda.get_device_name(0)
+    vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+    print(f"[env] gpu={name} vram={vram:.0f}GB bf16={bf16} torch={torch.__version__}",
+          flush=True)
+
+    model, processor = FastLanguageModel.from_pretrained(
+        model_name=MODEL_ID,
+        max_seq_length=args.max_seq,
+        dtype=None,
+        load_in_4bit=True,
+    )
+    tokenizer = get_tokenizer(processor)
+    print(f"[env] processor={type(processor).__name__} "
+          f"tokenizer={type(tokenizer).__name__}", flush=True)
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=32,
+        lora_alpha=32,
+        lora_dropout=0.0,
+        bias="none",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                        "gate_proj", "up_proj", "down_proj"],
+        use_gradient_checkpointing=False if args.no_grad_checkpoint else "unsloth",
+        random_state=SEED,
+        use_rslora=False,
+    )
+    model.print_trainable_parameters()
+
+    # The corpus is text-only, so no adapter may land on the vision tower. The
+    # projection names are shared across both stacks, so assert the scoping
+    # rather than trusting that the tower happens to be named differently.
+    adapted = {n.split(".lora_")[0] for n, _ in model.named_parameters() if ".lora_" in n}
+    stray = sorted(n for n in adapted if "language_model" not in n)
+    if stray:
+        raise SystemExit(f"LoRA landed outside the text decoder: {stray[:5]}")
+    print(f"[lora] {len(adapted)} modules adapted, all under language_model", flush=True)
+
+    data_dir = Path(args.data_dir)
+    train_ds = build_dataset(tokenizer, data_dir / "train.parquet",
+                             args.limit, args.sample)
+    eval_ds = build_dataset(tokenizer, data_dir / "validation.parquet",
+                            min(args.limit, 200) if args.limit else 0)
+    print(f"[data] train={len(train_ds)} validation={len(eval_ds)}", flush=True)
+
+    def n_tokens(text):
+        encoded = tokenizer(text)
+        ids = encoded["input_ids"] if hasattr(encoded, "keys") else encoded
+        return len(ids)
+
+    lengths = sorted(n_tokens(t) for t in train_ds["text"][:2000])
+    over = sum(1 for n in lengths if n > args.max_seq)
+    print(f"[data] tokens p50={lengths[len(lengths)//2]} "
+          f"p90={lengths[int(len(lengths)*0.9)]} max={lengths[-1]} "
+          f"over_max_seq={over}/{len(lengths)}", flush=True)
+
+    if args.check:
+        print("[check] sample render:\n" + train_ds["text"][0][:800], flush=True)
+        print("[check] done, not training", flush=True)
+        return
+
+    # trl renamed max_seq_length -> max_length and tokenizer -> processing_class
+    # between 0.11 and 1.x, and Colab resolves whatever is current. Pick by
+    # signature so a version bump is not another lost run.
+    import inspect
+
+    cfg_params = set(inspect.signature(SFTConfig.__init__).parameters)
+    seq_key = "max_length" if "max_length" in cfg_params else "max_seq_length"
+    trainer_params = set(inspect.signature(SFTTrainer.__init__).parameters)
+    tok_key = ("processing_class" if "processing_class" in trainer_params
+               else "tokenizer")
+    print(f"[api] trl seq_key={seq_key} tokenizer_key={tok_key}", flush=True)
+
+    cfg = SFTConfig(
+        **{seq_key: args.max_seq},
+        output_dir=args.output_dir,
+        per_device_train_batch_size=args.batch,
+        gradient_accumulation_steps=args.accum,
+        warmup_ratio=0.03,
+        num_train_epochs=args.epochs,
+        learning_rate=args.lr,
+        logging_steps=20,
+        optim="adamw_8bit",
+        weight_decay=0.01,
+        lr_scheduler_type="cosine",
+        seed=SEED,
+        bf16=bf16,
+        fp16=not bf16,
+        packing=False,          # conversations must stay intact
+        eval_strategy="steps",
+        eval_steps=args.eval_steps,
+        save_strategy="steps",
+        save_steps=args.save_steps,
+        save_total_limit=2,
+        report_to="none",
+        dataset_num_proc=2,
+        dataset_text_field="text",
+        group_by_length=args.group_by_length,
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        args=cfg,
+        **{tok_key: tokenizer},
+    )
+    trainer = train_on_responses_only(
+        trainer,
+        instruction_part=INSTRUCTION_PART,
+        response_part=RESPONSE_PART,
+    )
+
+    example = trainer.train_dataset[0]
+    trained = [t for t in example["labels"] if t != -100]
+    print(f"[mask] total={len(example['input_ids'])} trained={len(trained)} "
+          f"({100*len(trained)/max(1,len(example['input_ids'])):.1f}%)", flush=True)
+    if not trained:
+        raise SystemExit("label mask is empty -- the response delimiter did not match")
+
+    resume = newest_checkpoint(args.output_dir) if args.resume else None
+    print(f"[train] starting{f' from {resume}' if resume else ''}", flush=True)
+
+    start = time.time()
+    stats = trainer.train(resume_from_checkpoint=resume)
+    print(f"[train] finished in {(time.time()-start)/60:.1f} min", flush=True)
+    print(f"[train] metrics {stats.metrics}", flush=True)
+
+    os.makedirs(args.adapter_dir, exist_ok=True)
+    model.save_pretrained(args.adapter_dir)
+    # Save the processor, not just the inner tokenizer: reloading the adapter
+    # against a VL base needs the preprocessor config to be present.
+    processor.save_pretrained(args.adapter_dir)
+    print(f"[save] adapter -> {args.adapter_dir}", flush=True)
+    print("[done] TRAINING_COMPLETE", flush=True)
+
+
+if __name__ == "__main__":
+    main()

training/train_sahabatai.py

@@ -0,0 +1,267 @@
+"""QLoRA finetune of Sahabat-AI (Llama-3 8B) on the Fairleap chat corpus.
+
+Runs on the VM. Launch detached so the job outlives the websocket, the proxy
+token and the local session record:
+
+    nohup python -u train_sahabatai.py > /content/logs/train.log 2>&1 &
+
+The data pipeline, sampling and checkpoint handling are shared with
+`train_qwen.py` -- they are base-model agnostic. What is *not* shared, and what
+this file exists for, is everything Llama-3 gets wrong on this corpus:
+
+1. **The stock chat template silently destroys the tool data.** Sahabat-AI ships
+   the plain Llama-3 template: it renders `content | trim` for every message and
+   ignores `tool_calls` and `tools` entirely. An assistant tool-call turn carries
+   `content: ""`, so ~10% of the corpus would render as blank assistant replies
+   with no error anywhere. `fairleap_data/templates/llama3_tools.jinja` replaces
+   it and is verified off-GPU by `check_template.py`.
+2. **`model_max_length` is 2048** in the tokenizer config while the model itself
+   has 8192 positions. Left alone it truncates over a third of the corpus.
+3. **`generation_config.json` says eos is `<|end_of_text|>` (128001)** while the
+   chat template ends every turn with `<|eot_id|>` (128009). Training does not
+   care; inference does, and never stops without the override.
+"""
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from train_qwen import (SEED, build_dataset, get_tokenizer, newest_checkpoint,
+                        stratified_sample)  # noqa: F401  (re-exported for parity)
+
+MODEL_ID = "GoToCompany/llama3-8b-cpt-sahabatai-v1-instruct"
+
+TEMPLATE = Path(__file__).resolve().parent / "llama3_tools.jinja"
+
+# Llama-3 header delimiters. Tool results are rendered as `user` turns wrapping
+# <tool_response> precisely so this instruction delimiter matches them and they
+# stay out of the loss -- see the template.
+INSTRUCTION_PART = "<|start_header_id|>user<|end_header_id|>\n\n"
+RESPONSE_PART = "<|start_header_id|>assistant<|end_header_id|>\n\n"
+
+EOT_TOKEN = "<|eot_id|>"
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data-dir", default="/content/data")
+    p.add_argument("--output-dir", default="/content/outputs/sahabatai_8b_fairleap")
+    p.add_argument("--adapter-dir", default="/content/fairleap-v1-clm-sahabatai-8b-adapter")
+    p.add_argument("--template", default=str(TEMPLATE))
+    p.add_argument("--epochs", type=float, default=2.0)
+    p.add_argument("--max-seq", type=int, default=4096)
+    p.add_argument("--batch", type=int, default=2)
+    p.add_argument("--accum", type=int, default=8)
+    p.add_argument("--lr", type=float, default=2e-4)
+    p.add_argument("--save-steps", type=int, default=150)
+    p.add_argument("--eval-steps", type=int, default=150)
+    p.add_argument("--limit", type=int, default=0, help="cap records, for smoke runs")
+    p.add_argument("--sample", type=int, default=0,
+                   help="stratified subsample of train, preserving the mix")
+    p.add_argument("--max-steps", type=int, default=0, help="stop early, for smoke runs")
+    p.add_argument("--group-by-length", action="store_true",
+                   help="batch similar lengths together to cut padding waste")
+    p.add_argument("--no-grad-checkpoint", action="store_true",
+                   help="trade VRAM for speed; only if the card has headroom")
+    p.add_argument("--resume", action="store_true", help="continue from newest checkpoint")
+    p.add_argument("--check", action="store_true", help="render and report, do not train")
+    return p.parse_args()
+
+
+def install_template(tokenizer, path, max_seq):
+    """Replace the stock template and undo the two config traps."""
+    tokenizer.chat_template = Path(path).read_text(encoding="utf-8")
+
+    # 2048 in tokenizer_config.json against 8192 real positions. SFTTrainer
+    # honours model_max_length, so leaving it truncates most conversations.
+    if tokenizer.model_max_length < max_seq:
+        print(f"[fix] model_max_length {tokenizer.model_max_length} -> {max_seq}", flush=True)
+        tokenizer.model_max_length = max_seq
+
+    # Config ships padding_side=left, which is the inference setting.
+    tokenizer.padding_side = "right"
+
+    eot = tokenizer.convert_tokens_to_ids(EOT_TOKEN)
+    print(f"[fix] eos={tokenizer.eos_token!r} id={tokenizer.eos_token_id} "
+          f"eot_id={eot} pad={tokenizer.pad_token!r}", flush=True)
+    if tokenizer.eos_token_id != eot:
+        raise SystemExit(
+            f"eos is {tokenizer.eos_token!r}, not {EOT_TOKEN} -- the template ends "
+            f"turns with {EOT_TOKEN} and the model would never learn to stop")
+    return tokenizer
+
+
+def assert_text_only_lora(model):
+    """Every adapted module must be a decoder-layer projection.
+
+    Cheap insurance against a target list that quietly picks up an embedding or
+    an lm_head copy, which would multiply the adapter size and change what
+    "adapter-only" means for the repo.
+    """
+    adapted = {n.split(".lora_")[0] for n, _ in model.named_parameters() if ".lora_" in n}
+    stray = sorted(n for n in adapted if ".layers." not in n)
+    if stray:
+        raise SystemExit(f"LoRA landed outside the decoder layers: {stray[:5]}")
+    print(f"[lora] {len(adapted)} modules adapted, all decoder projections", flush=True)
+
+
+def main():
+    args = parse_args()
+
+    import torch
+    from unsloth import FastLanguageModel
+    from unsloth.chat_templates import train_on_responses_only
+    from trl import SFTConfig, SFTTrainer
+
+    bf16 = torch.cuda.is_bf16_supported()
+    name = torch.cuda.get_device_name(0)
+    vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+    print(f"[env] gpu={name} vram={vram:.0f}GB bf16={bf16} torch={torch.__version__}",
+          flush=True)
+
+    model, processor = FastLanguageModel.from_pretrained(
+        model_name=MODEL_ID,
+        max_seq_length=args.max_seq,
+        dtype=None,
+        load_in_4bit=True,
+    )
+    tokenizer = install_template(get_tokenizer(processor), args.template, args.max_seq)
+    print(f"[env] tokenizer={type(tokenizer).__name__} vocab={len(tokenizer)}", flush=True)
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=32,
+        lora_alpha=32,
+        lora_dropout=0.0,
+        bias="none",
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                        "gate_proj", "up_proj", "down_proj"],
+        use_gradient_checkpointing=False if args.no_grad_checkpoint else "unsloth",
+        random_state=SEED,
+        use_rslora=False,
+    )
+    model.print_trainable_parameters()
+    assert_text_only_lora(model)
+
+    data_dir = Path(args.data_dir)
+    train_ds = build_dataset(tokenizer, data_dir / "train.parquet",
+                             args.limit, args.sample)
+    eval_ds = build_dataset(tokenizer, data_dir / "validation.parquet",
+                            min(args.limit, 200) if args.limit else 0)
+    print(f"[data] train={len(train_ds)} validation={len(eval_ds)}", flush=True)
+
+    def n_tokens(text):
+        encoded = tokenizer(text)
+        ids = encoded["input_ids"] if hasattr(encoded, "keys") else encoded
+        return len(ids)
+
+    lengths = sorted(n_tokens(t) for t in train_ds["text"][:2000])
+    over = sum(1 for n in lengths if n > args.max_seq)
+    print(f"[data] tokens p50={lengths[len(lengths)//2]} "
+          f"p90={lengths[int(len(lengths)*0.9)]} max={lengths[-1]} "
+          f"over_max_seq={over}/{len(lengths)}", flush=True)
+
+    # The template is the single highest-risk piece: it fails by producing
+    # plausible text, not by raising. Prove a tool call survived the render.
+    rendered = "\n".join(train_ds["text"][:2000])
+    emitted = rendered.count('{"name": "predict_earnings"')
+    offered = rendered.count("Fungsi yang tersedia:")
+    results = rendered.count("<tool_response>")
+    print(f"[template] tool calls rendered={emitted} schemas offered={offered} "
+          f"tool results={results}", flush=True)
+    if offered and not emitted:
+        raise SystemExit("tools offered but no call rendered -- the template dropped them")
+
+    if args.check:
+        print("[check] sample render:\n" + train_ds["text"][0][:1200], flush=True)
+        print("[check] done, not training", flush=True)
+        return
+
+    # trl renamed max_seq_length -> max_length and tokenizer -> processing_class
+    # between 0.11 and 1.x, and Colab resolves whatever is current. Pick by
+    # signature so a version bump is not another lost run.
+    import inspect
+
+    cfg_params = set(inspect.signature(SFTConfig.__init__).parameters)
+    seq_key = "max_length" if "max_length" in cfg_params else "max_seq_length"
+    trainer_params = set(inspect.signature(SFTTrainer.__init__).parameters)
+    tok_key = ("processing_class" if "processing_class" in trainer_params
+               else "tokenizer")
+    print(f"[api] trl seq_key={seq_key} tokenizer_key={tok_key}", flush=True)
+
+    cfg = SFTConfig(
+        **{seq_key: args.max_seq},
+        output_dir=args.output_dir,
+        per_device_train_batch_size=args.batch,
+        gradient_accumulation_steps=args.accum,
+        warmup_ratio=0.03,
+        num_train_epochs=args.epochs,
+        max_steps=args.max_steps or -1,
+        learning_rate=args.lr,
+        logging_steps=10,
+        optim="adamw_8bit",
+        weight_decay=0.01,
+        lr_scheduler_type="cosine",
+        seed=SEED,
+        bf16=bf16,
+        fp16=not bf16,
+        packing=False,          # conversations must stay intact
+        eval_strategy="steps",
+        eval_steps=args.eval_steps,
+        save_strategy="steps",
+        save_steps=args.save_steps,
+        save_total_limit=2,
+        report_to="none",
+        dataset_num_proc=2,
+        dataset_text_field="text",
+        group_by_length=args.group_by_length,
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        args=cfg,
+        **{tok_key: tokenizer},
+    )
+    trainer = train_on_responses_only(
+        trainer,
+        instruction_part=INSTRUCTION_PART,
+        response_part=RESPONSE_PART,
+    )
+
+    example = trainer.train_dataset[0]
+    trained = [t for t in example["labels"] if t != -100]
+    print(f"[mask] total={len(example['input_ids'])} trained={len(trained)} "
+          f"({100*len(trained)/max(1,len(example['input_ids'])):.1f}%)", flush=True)
+    if not trained:
+        raise SystemExit("label mask is empty -- the response delimiter did not match")
+
+    # A tool result inside the loss would teach the model to write its own
+    # forecasts. Check the masked labels of a tool-bearing example, not the text.
+    for row in trainer.train_dataset:
+        kept = tokenizer.decode([t for t in row["labels"] if t != -100])
+        if "predict_earnings" in kept:
+            if "<tool_response>" in kept:
+                raise SystemExit("tool results are inside the loss -- masking is wrong")
+            print("[mask] tool call trained, tool result masked out", flush=True)
+            break
+
+    resume = newest_checkpoint(args.output_dir) if args.resume else None
+    print(f"[train] starting{f' from {resume}' if resume else ''}", flush=True)
+
+    start = time.time()
+    stats = trainer.train(resume_from_checkpoint=resume)
+    print(f"[train] finished in {(time.time()-start)/60:.1f} min", flush=True)
+    print(f"[train] metrics {stats.metrics}", flush=True)
+
+    os.makedirs(args.adapter_dir, exist_ok=True)
+    model.save_pretrained(args.adapter_dir)
+    tokenizer.save_pretrained(args.adapter_dir)
+    print(f"[save] adapter -> {args.adapter_dir}", flush=True)
+    print("[done] TRAINING_COMPLETE", flush=True)
+
+
+if __name__ == "__main__":
+    main()