"""Weight extrapolation (ExPO): theta = theta_inst + alpha * (theta_inst - theta_base)."""
import os, shutil, sys, torch
from safetensors.torch import load_file, save_file
alpha, out = float(sys.argv[1]), sys.argv[2]
inst = load_file("models/LFM3-1.7B/model.safetensors")
base = load_file("models/LFM3-1.7B-Base/model.safetensors")
new = {}
for k, v in inst.items():
    b = base.get(k)
    new[k] = v if b is None or b.shape != v.shape else (v.float() + alpha * (v.float() - b.float())).to(v.dtype)
os.makedirs(out, exist_ok=True)
save_file(new, os.path.join(out, "model.safetensors"), metadata={"format": "pt"})
for f in os.listdir("models/LFM3-1.7B"):
    if f != "model.safetensors" and os.path.isfile(os.path.join("models/LFM3-1.7B", f)):
        shutil.copy(os.path.join("models/LFM3-1.7B", f), out)
print("missing in base:", [k for k in inst if k not in base][:5])