Spaces:
Sleeping
Sleeping
File size: 3,816 Bytes
24e1932 0bf8fb1 24e1932 0bf8fb1 24e1932 a81fd74 24e1932 a07e79a 24e1932 e9f64da fa476b7 24e1932 0bf8fb1 24e1932 a07e79a e9f64da fa476b7 24e1932 a07e79a e9f64da fa476b7 24e1932 2dcdcdf 24e1932 0bf8fb1 7b59fb4 24e1932 | 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 100 101 102 103 104 105 106 | #!/usr/bin/env python3
"""
Primary pipeline stage: download model → extract weights → compute stats → HF Dataset.
Logs parameters, timing, and artifact path to MLflow.
Examples:
python scripts/run_primary.py --model gpt2 --out-dir outputs
python scripts/run_primary.py --model pythia-70m-deduped --revision step8000 --out-dir outputs
"""
import argparse
import os
import time
import hashlib
import json
import mlflow
from mlflow.entities import Dataset as DatasetEntity, DatasetInput
from transformer_analysis.head_pipeline import process_model
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--model", required=True)
p.add_argument("--revision", default="")
p.add_argument("--out-dir", default="outputs")
p.add_argument("--cache-dir", default="./model_data")
p.add_argument("--max-workers", type=int, default=4)
p.add_argument("--device", default=None)
p.add_argument("--weight-types", nargs="+", default=None)
p.add_argument("--cleanup-downloads", action="store_true")
p.add_argument("--low-rank-svd", action="store_true", default=False, dest="low_rank_svd")
p.add_argument("--top-k-svd", type=int, default=-1, dest="top_k_svd")
p.add_argument("--top-k-svd-d-head", action="store_true", default=False, dest="top_k_svd_d_head")
p.add_argument("--svd-via-gram", action="store_true", default=False, dest="svd_via_gram")
p.add_argument("--mlflow-uri", default="file:./mlruns")
p.add_argument("--mlflow-experiment", default="production")
return p.parse_args()
def main():
args = parse_args()
revision = args.revision or None
mlflow.set_tracking_uri(args.mlflow_uri)
mlflow.set_experiment(args.mlflow_experiment)
rev_label = revision or "main"
run_name = f"{args.model}-{rev_label}-primary"
with mlflow.start_run(run_name=run_name):
mlflow.set_tag("mlflow.note.content", "Extract per-head statistics from model weights → HF Dataset")
mlflow.set_tag("model", args.model)
mlflow.log_params({
"model": args.model,
"revision": rev_label,
"max_workers": args.max_workers,
"device": args.device or "auto",
"out_dir": args.out_dir,
"weight_types": ",".join(args.weight_types) if args.weight_types else "default",
"low_rank_svd": args.low_rank_svd,
"top_k_svd": args.top_k_svd,
"top_k_svd_d_head": args.top_k_svd_d_head,
"svd_via_gram": args.svd_via_gram,
})
t0 = time.time()
process_model(
model_name=args.model,
revision=revision,
out_dir=args.out_dir,
cache_dir=args.cache_dir,
cleanup_downloads=args.cleanup_downloads,
max_workers=args.max_workers,
device=args.device,
weight_types=args.weight_types,
low_rank_svd_approximation=args.low_rank_svd,
top_k_svd=args.top_k_svd,
top_k_svd_d_head=args.top_k_svd_d_head,
svd_via_gram=args.svd_via_gram,
)
wall_time = time.time() - t0
out_key = f"{args.model}_{revision or 'main'}"
dataset_path = os.path.join(args.out_dir, out_key)
mlflow.log_metric("wall_time_s", wall_time)
mlflow.log_param("dataset_path", dataset_path)
digest = hashlib.md5(dataset_path.encode()).hexdigest()[:8]
entity = DatasetEntity(
name=f"{args.model}-{rev_label}",
digest=digest,
source_type="local",
source=json.dumps({"uri": dataset_path}),
)
mlflow.MlflowClient().log_inputs(
mlflow.active_run().info.run_id,
[DatasetInput(dataset=entity, tags=[])],
)
if __name__ == "__main__":
main()
|