The current best open-weight embedding model.

Pull, serve, truncate, and normalize.

import argparse
import json
import shutil
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file

COPY_FILES = (
    "config.json",
    "generation_config.json",
    "tokenizer.json",
    "tokenizer.model",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "preprocessor_config.json",
    "processor_config.json",
    "chat_template.jinja",
    "chat_template.json",
)

def weight_map(model_dir: Path) -> dict[str, Path]:
    index = model_dir / "model.safetensors.index.json"
    if index.is_file():
        meta = json.loads(index.read_text())
        return {k: model_dir / v for k, v in meta["weight_map"].items()}
    single = model_dir / "model.safetensors"
    if not single.is_file():
        raise FileNotFoundError(f"no safetensors found in {model_dir}")
    with safe_open(str(single), framework="pt") as f:
        return {k: single for k in f.keys()}

def load(path: Path, key: str) -> torch.Tensor:
    with safe_open(str(path), framework="pt") as f:
        return f.get_tensor(key)

def select_keys(wm: dict[str, Path], key_cutoff: str = None, layer_index: int = None) -> list[str]:
    keys = list(wm.keys())

    cutoff_pattern = None
    if layer_index is not None:
        cutoff_pattern = f"layers.{layer_index}."
    elif key_cutoff is not None:
        cutoff_pattern = key_cutoff

    if cutoff_pattern:
        matching = [k for k in keys if cutoff_pattern in k]
        if not matching:
            raise KeyError(f"No keys matching pattern '{cutoff_pattern}' found in model.")
        stop_key = matching[0]
        idx = keys.index(stop_key)
        keys = keys[:idx + 1]

    return keys

def process_sidecars(src: Path, dst: Path, keys: list[str]):
    max_layer = -1
    for k in keys:
        if "layers" in k:
            parts = k.split(".")
            try:
                idx = parts.index("layers")
                layer_num = int(parts[idx + 1])
                if layer_num > max_layer:
                    max_layer = layer_num
            except (ValueError, IndexError):
                pass

    for name in COPY_FILES:
        p = src / name
        if not p.is_file():
            continue

        if name.endswith(".index.json"):
            continue

        if name == "config.json":
            try:
                config_data = json.loads(p.read_text())
                if max_layer >= 0 and "num_hidden_layers" in config_data:
                    new_layers = max_layer + 1
                    config_data["num_hidden_layers"] = new_layers
                    print(f"Updated config.json: num_hidden_layers -> {new_layers}")

                (dst / name).write_text(json.dumps(config_data, indent=2))
                print(f"processed and saved {name}")
                continue
            except Exception as e:
                print(f"Warning: Could not update config.json ({e}), copying as-is.")

        shutil.copy2(p, dst / name)
        print(f"copied {name}")

def main():
    parser = argparse.ArgumentParser(description="Extract sequential weight slices from a Hugging Face model.")
    parser.add_argument("--model", required=True, help="Path to model directory")
    parser.add_argument("--key", default=None, help="Cut off extraction at the first key matching this substring")
    parser.add_argument("--index", type=int, default=None, help="Cut off extraction up to this transformer layer index (e.g. 0, 5, 11)")
    parser.add_argument("--print", action="store_true", help="Print all keys and exit")
    args = parser.parse_args()

    model_dir = Path(args.model).resolve()
    wm = weight_map(model_dir)

    if args.print:
        keys = select_keys(wm, key_cutoff=args.key, layer_index=args.index)
        print(f"Found {len(keys)} keys up to cutoff:")
        for k in keys:
            print(k)
        return

    # Determine safe output suffix name
    if args.index is not None:
        suffix = f"layer_{args.index}"
    else:
        suffix = args.key or "all"

    safe_suffix = "".join(c if c.isalnum() else "_" for c in suffix)[:30]
    out_dir = Path(f"{model_dir.name}-extracted-{safe_suffix}")
    out_dir.mkdir(parents=True, exist_ok=True)

    keys = select_keys(wm, key_cutoff=args.key, layer_index=args.index)

    tensors = {}
    for k in keys:
        t = load(wm[k], k).contiguous()
        tensors[k] = t
        print(f"{k:72s} {str(tuple(t.shape)):20s} {wm[k].name}")

    out = out_dir / "model.safetensors"
    save_file(tensors, str(out))

    process_sidecars(model_dir, out_dir, keys)
    print(f"wrote {len(tensors)} tensors -> {out}")

if __name__ == "__main__":
    main()