Spaces:
Sleeping
Sleeping
File size: 5,516 Bytes
a81fd74 f49060a a81fd74 dbd0f5c a81fd74 0bf8fb1 dbd0f5c 0bf8fb1 f49060a a81fd74 f49060a a81fd74 f49060a a81fd74 f49060a a81fd74 937097a 66a0324 dbd0f5c a81fd74 dbd0f5c a81fd74 dbd0f5c 0bf8fb1 66a0324 dbd0f5c 937097a 66a0324 0bf8fb1 dbd0f5c 937097a 66a0324 dbd0f5c 0bf8fb1 88dbab4 a81fd74 | 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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | #!/usr/bin/env python3
"""Thin CLI wrapper for eval_metrics. Called by rules/eval.smk."""
import argparse
import os
import sys
import time
from pathlib import Path
import hashlib
import json
import mlflow
from mlflow.entities import Dataset as DatasetEntity, DatasetInput
import pandas as pd
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
from transformer_analysis.eval_metrics import evaluate_model, to_long_format
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--revision", type=str, default=None)
parser.add_argument("--corpus", type=str, default="wikitext103",
choices=["wikitext103", "pile"])
parser.add_argument("--pile-tokens", type=int, default=204800)
parser.add_argument("--pile-cache", type=str, default=None)
parser.add_argument("--max-tokens", type=int, default=None)
parser.add_argument("--stride", type=int, default=512)
parser.add_argument("--out", type=str, required=True)
parser.add_argument("--cache", type=str, default="model_cache")
parser.add_argument("--device", type=str, default=None, choices=["cuda", "mps", "cpu"])
parser.add_argument("--dtype", type=str, default="auto",
choices=["auto", "bf16", "fp16", "fp32"],
help="Weight dtype. 'auto' uses the model's saved dtype (bf16/fp16 for most modern models, fp32 for GPT-2). Override with fp32 for a numerics reference.")
parser.add_argument("--offload-layers", action="store_true",
help="Stream layers from disk via accelerate. Enables models that don't fit fully in unified memory at the cost of much slower per-window inference. Numerics are unchanged.")
parser.add_argument("--offload-folder", type=str, default=None,
help="Where accelerate spills offloaded layer weights. Defaults to a tempfile.mkdtemp() on the system disk (best per-window bandwidth, auto-cleaned). Override to e.g. an external drive when the system disk is tight on space — pay 2-5x per-window cost in exchange.")
parser.add_argument("--max-memory-gpu", type=str, default=None,
help="Max memory budget for the accelerator (CUDA/MPS) when --offload-layers is on, e.g. '6GiB'. Layers beyond this spill to CPU then disk.")
parser.add_argument("--max-memory-cpu", type=str, default=None,
help="Max memory budget for CPU when --offload-layers is on, e.g. '2GiB'.")
parser.add_argument("--mlflow-uri", default="file:./mlruns")
parser.add_argument("--mlflow-experiment", default="production")
args = parser.parse_args()
revision = args.revision or None
rev_label = revision or "main"
mlflow.set_tracking_uri(args.mlflow_uri)
mlflow.set_experiment(args.mlflow_experiment)
with mlflow.start_run(run_name=f"{args.model}-{rev_label}-eval"):
mlflow.set_tag("mlflow.note.content", "Evaluate model perplexity on corpus → parquet")
mlflow.set_tag("model", args.model)
# Assemble accelerate max_memory dict only if either knob is set.
max_memory = None
if args.offload_layers and (args.max_memory_gpu or args.max_memory_cpu):
max_memory = {}
if args.max_memory_gpu:
# Key resolved inside evaluate_model based on actual device.
max_memory["__gpu__"] = args.max_memory_gpu
if args.max_memory_cpu:
max_memory["cpu"] = args.max_memory_cpu
mlflow.log_params({
"model": args.model,
"revision": rev_label,
"corpus": args.corpus,
"pile_tokens": args.pile_tokens,
"stride": args.stride,
"dtype": args.dtype,
"device": args.device or "auto",
"offload_layers": args.offload_layers,
"max_memory_gpu": args.max_memory_gpu or "",
"max_memory_cpu": args.max_memory_cpu or "",
"out": args.out,
})
t0 = time.time()
result = evaluate_model(
model_name=args.model, revision=revision,
corpus=args.corpus, pile_tokens=args.pile_tokens,
cache_dir=args.cache, device_str=args.device,
stride=args.stride, max_tokens=args.max_tokens,
pile_cache=args.pile_cache, dtype=args.dtype,
offload=args.offload_layers,
offload_folder=args.offload_folder,
max_memory=max_memory,
)
wall_time = time.time() - t0
print(f" ppl={result['perplexity']:.2f} bpb={result['bpb']:.4f}")
mlflow.log_metrics({
"perplexity": result["perplexity"],
"nll": result["nll"],
"bpb": result["bpb"],
"wall_time_s": wall_time,
})
df = pd.DataFrame(to_long_format(result))
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
df.to_parquet(args.out, index=False)
print(f" Saved {len(df)} rows → {args.out}")
digest = hashlib.md5(args.out.encode()).hexdigest()[:8]
entity = DatasetEntity(
name=f"{args.model}-{rev_label}-eval",
digest=digest,
source_type="local",
source=json.dumps({"uri": args.out}),
)
mlflow.MlflowClient().log_inputs(
mlflow.active_run().info.run_id,
[DatasetInput(dataset=entity, tags=[])],
)
if __name__ == "__main__":
main()
|