File size: 1,771 Bytes
36c78b1 |
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 |
import argparse
import json
from pathlib import Path
import torch
from bit_transformer import BitTransformerLM
def list_candidates(path: Path):
models = sorted(path.glob("*.pt"))
for m in models:
metrics_file = m.with_suffix(m.suffix + ".json")
metrics = {}
if metrics_file.exists():
with open(metrics_file) as f:
metrics = json.load(f)
yield m, metrics
def main():
parser = argparse.ArgumentParser(description="Review distilled submodels")
parser.add_argument("directory", type=Path, help="Directory with candidate models")
parser.add_argument("--approve-dir", type=Path, default=Path("approved"), help="Directory to store approved models")
args = parser.parse_args()
args.approve_dir.mkdir(exist_ok=True)
log_file = args.approve_dir / "review_log.jsonl"
for model_path, metrics in list_candidates(args.directory):
print("Candidate:", model_path.name)
for k, v in metrics.items():
print(f" {k}: {v}")
ans = input("Approve this model? [y/N] ").strip().lower()
if ans == "y":
approved_path = args.approve_dir / model_path.name
torch.save(torch.load(model_path), approved_path)
entry = {"model": approved_path.name, "metrics": metrics, "approved": True}
with open(log_file, "a") as lf:
lf.write(json.dumps(entry) + "\n")
print("Approved and saved to", approved_path)
else:
entry = {"model": model_path.name, "metrics": metrics, "approved": False}
with open(log_file, "a") as lf:
lf.write(json.dumps(entry) + "\n")
print("Rejected", model_path.name)
if __name__ == "__main__":
main()
|