File size: 3,476 Bytes
d72e738 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Rebuild a full TSLM ``best_model.pt`` state dict from this repository.
Only the tensors training changed are stored here. Everything else is recovered
bitwise from the public base checkpoints (Qwen3.5-4B, Chronos-2), so the result
is byte-identical to the original training artifact.
python reconstruct.py --run anomaly_xl_itformer_precise_chronos2_qwen35_4b \
--out best_model.pt
In-process, without writing a ``.pt``::
from reconstruct import rebuild_state_dict
state = rebuild_state_dict("anomaly_xl_itformer_precise_chronos2_qwen35_4b")
model.load_state_dict(state)
"""
import argparse
import glob
import json
from pathlib import Path
import torch
from huggingface_hub import snapshot_download
from safetensors import safe_open
REPO_ID = "nz00shuuuu/TimeRLM-TSLM-Checkpoints"
DTYPES = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}
_STOCK_CACHE = {}
def load_safetensors(paths):
out = {}
for p in paths:
with safe_open(str(p), framework="pt") as f:
for k in f.keys():
out[k] = f.get_tensor(k)
return out
def load_stock():
"""Download and load the public base checkpoints, memoised per process."""
if not _STOCK_CACHE:
qdir = snapshot_download("Qwen/Qwen3.5-4B", allow_patterns=["*.safetensors"])
cdir = snapshot_download("amazon/chronos-2", allow_patterns=["*.safetensors"])
_STOCK_CACHE["stock_qwen"] = load_safetensors(sorted(glob.glob(f"{qdir}/*.safetensors")))
_STOCK_CACHE["stock_chronos"] = load_safetensors(sorted(glob.glob(f"{cdir}/*.safetensors")))
return _STOCK_CACHE
def run_dir(run=None, local_dir=None):
if local_dir:
return Path(local_dir)
if not run:
raise ValueError("pass --run or --local-dir")
root = snapshot_download(REPO_ID, allow_patterns=[f"tslm/{run}/*"])
return Path(root) / "tslm" / run
def rebuild_state_dict(run=None, local_dir=None, stock=None):
"""Return the full ``model_state``, bitwise-equal to the original checkpoint."""
root = run_dir(run, local_dir)
meta = json.loads((root / "manifest.json").read_text())
trained = load_safetensors([root / "trained.safetensors"])
stock = stock or load_stock()
state = {}
for key in meta["key_order"]:
e = meta["tensors"][key]
src = e["from"]
if src == "extract":
state[key] = trained[key]
elif src == "alias":
state[key] = state[e["key"]]
else:
t = stock[src][e["key"]].to(DTYPES[e["dtype"]])
if "slice_rows" in e:
t = t[: e["slice_rows"]].contiguous()
state[key] = t
return state
def main():
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument("--run", help="run name, e.g. anomaly_xl_itformer_precise_chronos2_qwen35_4b")
ap.add_argument("--local-dir", help="use an already-downloaded run directory")
ap.add_argument("--out", help="write a train.py-compatible .pt here")
args = ap.parse_args()
state = rebuild_state_dict(args.run, args.local_dir)
print(f"rebuilt {len(state)} tensors")
if args.out:
meta = json.loads((run_dir(args.run, args.local_dir) / "manifest.json").read_text())
torch.save({"model_state": state, "config": meta["config"], **meta["training"]}, args.out)
print(f"wrote {args.out}")
if __name__ == "__main__":
main()
|