spb/forge Public MIT

1# Author: Simon-Pierre Boucher — contact@spboucher.ai2"""Inspect and convert .forge weight repositories (see src/core/fmodel.h for3the format: git-style manifests + content-addressed page-aligned shards).45Usage:6  python3 tools/fmodel.py inspect  model.forge [--manifest manifest-000002.json]7  python3 tools/fmodel.py log      model.forge                 # manifest history8  python3 tools/fmodel.py to-safetensors model.forge out.safetensors9        [--manifest m.json]        # pure-numpy writer, no torch needed1011Only numpy is required. The safetensors export makes .forge weights loadable12from PyTorch/HF (`safetensors.torch.load_file`) for interop.13"""14import argparse15import json16import os17import struct1819import numpy as np2021PAGE = 1638422NP_DTYPE = {"f32": np.float32, "f16": np.float16, "u16": np.uint16, "i32": np.int32}23ST_DTYPE = {"f32": "F32", "f16": "F16", "bf16": "BF16", "u16": "U16", "i32": "I32"}242526def resolve_manifest(repo, manifest=None):27    if os.path.isfile(repo):  # given a manifest path directly28        return os.path.dirname(os.path.dirname(os.path.abspath(repo))) \29            if os.path.basename(os.path.dirname(repo)) == "manifests" \30            else os.path.dirname(os.path.abspath(repo)), repo31    path = os.path.join(repo, "manifests", manifest) if manifest else \32        os.path.join(repo, "manifest-latest.json")33    return repo, path343536def load(repo, manifest=None):37    repo, mpath = resolve_manifest(repo, manifest)38    with open(mpath) as f:39        return repo, json.load(f)404142def tensor_bytes(repo, entry):43    with open(os.path.join(repo, entry["shard"]), "rb") as f:44        f.seek(entry["offset"])45        return f.read(entry["nbytes"])464748def tensor_array(repo, entry):49    raw = tensor_bytes(repo, entry)50    dt = entry["dtype"]51    if dt == "bf16":  # numpy has no bf16: widen via bit tricks52        u16 = np.frombuffer(raw, dtype=np.uint16).astype(np.uint32) << 1653        return u16.view(np.float32).reshape(entry["shape"])54    return np.frombuffer(raw, dtype=NP_DTYPE[dt]).reshape(entry["shape"])555657def cmd_inspect(args):58    repo, m = load(args.repo, args.manifest)59    cfg = m.get("config", {}).get("model", {})60    print(f"{m.get('self', '?')}  step={m.get('step')}  tag='{m.get('tag', '')}'"61          f"  dtype={m.get('dtype')}  parent='{m.get('parent', '')}'")62    if cfg:63        print(f"model: {cfg.get('name')}  layers={cfg.get('n_layers')}"64              f"  d_model={cfg.get('d_model')}  vocab={cfg.get('vocab_size')}")65    total = 066    shards = {}67    for name, e in sorted(m["tensors"].items()):68        total += e["nbytes"]69        shards.setdefault(e["shard"], 0)70        shards[e["shard"]] += e["nbytes"]71        print(f"  {name:<40} {e['dtype']:<5} {str(e['shape']):<20}"72              f" {e['nbytes'] / 1e6:8.2f} MB  {e['shard'].split('/')[-1][:8]}…")73    print(f"{len(m['tensors'])} tensors, {total / 1e6:.1f} MB across "74          f"{len(shards)} shard(s)")757677def cmd_log(args):78    repo, _ = resolve_manifest(args.repo)79    mdir = os.path.join(repo, "manifests")80    for name in sorted(os.listdir(mdir)):81        with open(os.path.join(mdir, name)) as f:82            m = json.load(f)83        shards = {e["shard"] for e in m["tensors"].values()}84        print(f"{name}  step={m.get('step'):>8}  tag='{m.get('tag', '')}'"85              f"  dtype={m.get('dtype')}  shards={len(shards)}"86              f"  parent='{m.get('parent', '')}'")878889def cmd_to_safetensors(args):90    repo, m = load(args.repo, args.manifest)91    header = {"__metadata__": {"format": "forge",92                               "step": str(m.get("step", -1)),93                               "config": json.dumps(m.get("config", {}))}}94    offset = 095    order = sorted(m["tensors"].items())96    for name, e in order:97        header[name] = {"dtype": ST_DTYPE[e["dtype"]], "shape": e["shape"],98                        "data_offsets": [offset, offset + e["nbytes"]]}99        offset += e["nbytes"]100    hjson = json.dumps(header, separators=(",", ":")).encode()101    pad = (8 - len(hjson) % 8) % 8  # safetensors aligns the header to 8102    hjson += b" " * pad103    with open(args.out, "wb") as f:104        f.write(struct.pack("<Q", len(hjson)))105        f.write(hjson)106        for name, e in order:107            f.write(tensor_bytes(repo, e))108    print(f"wrote {args.out} ({(8 + len(hjson) + offset) / 1e6:.1f} MB, "109          f"{len(order)} tensors)")110111112def main():113    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])114    sub = ap.add_subparsers(dest="cmd", required=True)115    for name, fn in [("inspect", cmd_inspect), ("log", cmd_log),116                     ("to-safetensors", cmd_to_safetensors)]:117        p = sub.add_parser(name)118        p.add_argument("repo", help=".forge directory or manifest json")119        p.add_argument("--manifest", default=None,120                       help="specific manifest inside manifests/")121        if name == "to-safetensors":122            p.add_argument("out")123        p.set_defaults(fn=fn)124    args = ap.parse_args()125    args.fn(args)126127128if __name__ == "__main__":129    main()130