deploy/export_model.py

#!/usr/bin/env python3
"""Export a trained checkpoint out of Tinker for self-hosting.

    python3 deploy/export_model.py --checkpoint tinker://<run-id>/sampler_weights/final

This matters for the IP-privacy case: the weights come out, so inference can run
entirely on your own infrastructure and no customer disclosure ever leaves it.

Two output modes:

  --mode adapter  (default)  a PEFT LoRA adapter, a few hundred MB.
                             Serve with:
                               vllm serve Qwen/Qwen3.5-9B \\
                                 --lora-modules claim-drafter=./peft_adapter
                             Lets you hot-swap adapters without redeploying the base.

  --mode merged              the base model with the adapter merged in, ~18GB.
                             Use only if your serving stack cannot load LoRA.

Deliberately NOT wired up: tinker_cookbook.weights.publish_to_hf_hub. Pushing a
model trained on customer disclosures to a public hub is the exact opposite of
the privacy guarantee. If you ever want that, do it by hand and with intent.
"""

import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from claim_drafter.config import DEFAULTS, load_env


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", help="tinker:// path (omit if --adapter-path is given)")
    ap.add_argument("--adapter-path", help="reuse an already-downloaded raw adapter "
                    "dir (skips the Tinker download entirely)")
    ap.add_argument("--base-model", default=DEFAULTS["model"])
    ap.add_argument("--mode", choices=["adapter", "merged"], default="adapter")
    ap.add_argument("--work-dir", default="export")
    args = ap.parse_args()

    if not args.checkpoint and not args.adapter_path:
        ap.error("pass either --checkpoint (to download) or --adapter-path (to reuse)")

    load_env()
    # build_hf_model loads the base model's config from the Hub. Qwen configs
    # need remote code, and resolve_trust_remote_code() defaults to False, so
    # --mode merged dies immediately without this.
    os.environ.setdefault("HF_TRUST_REMOTE_CODE", "1")
    from tinker_cookbook import weights

    os.makedirs(args.work_dir, exist_ok=True)
    if args.adapter_path:
        adapter_dir = args.adapter_path
        print("Reusing adapter at %s (no download)" % adapter_dir)
    else:
        raw_dir = os.path.join(args.work_dir, "raw_adapter")
        print("Downloading %s -> %s" % (args.checkpoint, raw_dir))
        adapter_dir = weights.download(tinker_path=args.checkpoint, output_dir=raw_dir)

    if args.mode == "adapter":
        out = os.path.join(args.work_dir, "peft_adapter")
        weights.build_lora_adapter(
            base_model=args.base_model, adapter_path=adapter_dir, output_path=out)
        print("\nPEFT adapter written to %s" % out)
        print("\nServe it with:")
        print("  vllm serve %s \\\n    --lora-modules claim-drafter=%s \\\n"
              "    --max-model-len 16384 --port 8000" % (args.base_model, out))
    else:
        out = os.path.join(args.work_dir, "model")
        weights.build_hf_model(
            base_model=args.base_model, adapter_path=adapter_dir, output_path=out)
        print("\nMerged model written to %s" % out)
        print("\nServe it with:")
        print("  vllm serve %s --max-model-len 16384 --port 8000" % out)


if __name__ == "__main__":
    main()