#!/usr/bin/env python3
from pathlib import Path

import torch


ckpt = Path(
    "runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/"
    "checkpoints/iter_0107500"
)

print("metadata.json:", (ckpt / "metadata.json").read_text())
common = torch.load(ckpt / "common.pt", map_location="cpu")
print("common iteration:", common.get("iteration"))
print("common keys:", list(common.keys()))
content_metadata = common.get("content_metadata")
print("content_metadata type:", type(content_metadata))
print("content_metadata:", content_metadata)

print("run_config relevant lines:")
for line in (ckpt / "run_config.yaml").read_text().splitlines():
    if any(
        needle in line.lower()
        for needle in ["hf", "hugging", "auto", "architecture", "tokenizer", "model_provider", "bridge"]
    ):
        print(line[:240])