| |
| """Average two of our own checkpoints (same architecture).""" |
| from __future__ import annotations |
|
|
| import json |
| import shutil |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file, save_file |
|
|
|
|
| def main(): |
| a, b, out, wa = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3]), float(sys.argv[4]) |
| wb = 1.0 - wa |
| out.mkdir(parents=True, exist_ok=True) |
| shards = sorted(a.glob("model-*.safetensors")) |
| if not shards: |
| raise SystemExit(f"no shards in {a}") |
| index = {"metadata": {"format": "pt"}, "weight_map": {}} |
| for sa in shards: |
| sb = b / sa.name |
| ta = load_file(str(sa), device="cpu") |
| tb = load_file(str(sb), device="cpu") |
| merged = {} |
| for k in ta: |
| if k not in tb: |
| merged[k] = ta[k] |
| continue |
| if ta[k].dtype in (torch.float16, torch.bfloat16, torch.float32): |
| merged[k] = (ta[k].float() * wa + tb[k].float() * wb).to(ta[k].dtype) |
| else: |
| merged[k] = ta[k] |
| dest = out / sa.name |
| save_file(merged, str(dest)) |
| for k in merged: |
| index["weight_map"][k] = sa.name |
| print("wrote", dest.name, "tensors", len(merged)) |
| (out / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) |
| for name in ( |
| "config.json", |
| "generation_config.json", |
| "tokenizer.json", |
| "tokenizer_config.json", |
| "chat_template.jinja", |
| "special_tokens_map.json", |
| ): |
| src = a / name |
| if src.exists(): |
| shutil.copy2(src, out / name) |
| (out / "STABLE").write_text("ok\n") |
| print("done", out, "w", wa, wb) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
