299 lines
17 KiB
Python
299 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
"""scan_experiments.py - inventory of the machine-learning experiments in a project (stdlib only).
|
|
|
|
python scan_experiments.py [ROOT] [--out paper/inventory]
|
|
|
|
Walks ROOT (default '.'), skipping paper/, .git, venvs, node_modules, caches, and classifies:
|
|
code .py / .ipynb files that import an ML framework; entry points (__main__, argparse,
|
|
hydra, click), CLI arguments with defaults, seeds, datasets, models, optimizers, metrics
|
|
configs .yaml/.yml/.json/.toml files that look like hyperparameter configs
|
|
results .csv/.tsv/.json/.jsonl result tables, *.log files with metric lines, TensorBoard
|
|
event files, wandb/ and mlruns/ folders, checkpoints (only counted)
|
|
docs README*, *.md, requirements.txt, environment.yml, pyproject.toml (dependencies)
|
|
Writes <out>.json (full detail) and <out>.md (human-readable), prints a compact summary.
|
|
The script only reads; nothing is executed.
|
|
"""
|
|
import ast, json, re, sys
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
SKIP = {"paper", ".git", "__pycache__", ".venv", "venv", "env", "node_modules", ".ipynb_checkpoints",
|
|
".pytest_cache", ".mypy_cache", ".ruff_cache", "dist", "build", "site-packages", ".idea", ".vscode"}
|
|
ML_LIBS = {"torch": "PyTorch", "tensorflow": "TensorFlow", "keras": "Keras", "jax": "JAX", "flax": "Flax",
|
|
"sklearn": "scikit-learn", "xgboost": "XGBoost", "lightgbm": "LightGBM", "catboost": "CatBoost",
|
|
"transformers": "Hugging Face Transformers", "datasets": "Hugging Face Datasets",
|
|
"lightning": "PyTorch Lightning", "pytorch_lightning": "PyTorch Lightning",
|
|
"torchvision": "torchvision", "timm": "timm", "peft": "PEFT", "trl": "TRL", "accelerate": "Accelerate",
|
|
"optuna": "Optuna", "statsmodels": "statsmodels", "gymnasium": "Gymnasium", "gym": "Gym",
|
|
"stable_baselines3": "Stable-Baselines3", "torch_geometric": "PyG", "dgl": "DGL",
|
|
"sentence_transformers": "Sentence-Transformers", "wandb": "Weights & Biases", "mlflow": "MLflow",
|
|
"numpy": "NumPy", "pandas": "pandas", "scipy": "SciPy"}
|
|
CORE_ML = set(ML_LIBS) - {"numpy", "pandas", "scipy", "wandb", "mlflow", "optuna"}
|
|
DATASETS = r"""MNIST|Fashion-?MNIST|KMNIST|EMNIST|CIFAR-?10|CIFAR-?100|SVHN|STL-?10|ImageNet(?:-1k|-21k)?|Tiny-?ImageNet|
|
|
CelebA|LSUN|COCO|MS-?COCO|Pascal\s?VOC|Cityscapes|ADE20K|Places365|Caltech-?101|Caltech-?256|Oxford-?Pets|Flowers-?102|Food-?101|
|
|
IMDB|SST-?2|GLUE|SuperGLUE|SQuAD|MNLI|QNLI|QQP|RTE|CoLA|AG-?News|Yelp|Amazon|WikiText-?\d*|Penn\s?Treebank|PTB|OpenWebText|
|
|
C4|The\s?Pile|LAMBADA|HellaSwag|MMLU|GSM8K|HumanEval|MBPP|ARC|TruthfulQA|WMT\d*|Multi30k|LibriSpeech|Common\s?Voice|
|
|
UCI|Adult|Iris|Wine|Breast\s?Cancer|California\s?Housing|Boston|Covertype|Higgs|Titanic|OpenML|
|
|
Cora|CiteSeer|PubMed|OGB[\w-]*|QM9|ZINC|MovieLens|Criteo|ETTh\d|ETTm\d|Electricity|Traffic|Weather|M4|
|
|
CartPole|MountainCar|Atari|MuJoCo|HalfCheetah|Hopper|Walker2d|Ant|Humanoid|Procgen|D4RL"""
|
|
DATASET_RE = re.compile(r"(?<![\w-])(" + DATASETS.replace("\n", "") + r")(?![\w-])", re.I)
|
|
HF_DS_RE = re.compile(r"load_dataset\(\s*['\"]([^'\"]+)['\"](?:\s*,\s*['\"]([^'\"]+)['\"])?")
|
|
MODEL_RE = re.compile(r"(?<![\w])(ResNet-?\d+|ResNeXt-?\d*|WideResNet[\w-]*|WRN-?\d+-?\d*|VGG-?\d+|DenseNet-?\d+|"
|
|
r"MobileNet(?:V\d)?|EfficientNet-?\w*|ViT-?[\w/]*|DeiT[\w-]*|Swin[\w-]*|ConvNeXt[\w-]*|U-?Net|"
|
|
r"BERT[\w-]*|RoBERTa[\w-]*|DistilBERT|ALBERT|DeBERTa[\w-]*|ELECTRA|T5[\w-]*|BART|GPT-?\w*|"
|
|
r"Llama[\w.-]*|Mistral[\w-]*|Qwen[\w.-]*|Gemma[\w-]*|Phi-?\d[\w-]*|CLIP|LSTM|GRU|Transformer|"
|
|
r"MLP|CNN|GCN|GAT|GraphSAGE|GIN|VAE|GAN|DCGAN|StyleGAN\d?|DDPM|DDIM|Diffusion|"
|
|
r"RandomForest\w*|GradientBoosting\w*|LogisticRegression|SVC|SVM|KNeighbors\w*|XGB\w*|LGBM\w*|"
|
|
r"PPO|SAC|DQN|A2C|TD3|DDPG)(?![\w])", re.I)
|
|
OPTIM_RE = re.compile(r"\b(SGD|Adam|AdamW|RMSprop|Adagrad|Adafactor|LAMB|LARS|Lion)\b")
|
|
METRIC_WORDS = r"acc(?:uracy)?|top-?[15]|loss|f1|precision|recall|auc|auroc|auprc|mse|rmse|mae|r2|bleu|rouge[\w-]*|" \
|
|
r"perplexity|ppl|map|miou|iou|dice|psnr|ssim|fid|is_score|reward|return|ece|nll|wer|cer|em|exact_match"
|
|
METRIC_RE = re.compile(r"(?<![a-z])(" + METRIC_WORDS + r")(?![a-z])", re.I)
|
|
SEED_RE = re.compile(r"(?:manual_seed|random\.seed|np\.random\.seed|seed_everything|set_seed|random_state)\s*[(=]\s*([\w.]+)")
|
|
LOG_METRIC_RE = re.compile(r"(?i)\b(" + METRIC_WORDS + r"|val_\w+|test_\w+|train_\w+)\s*[:=]\s*(-?\d+\.?\d*(?:e-?\d+)?)")
|
|
HP_KEYS = {"lr", "learning_rate", "batch_size", "epochs", "num_epochs", "max_epochs", "weight_decay", "momentum",
|
|
"dropout", "hidden_size", "hidden_dim", "num_layers", "seed", "optimizer", "warmup", "max_steps",
|
|
"n_estimators", "max_depth", "model", "dataset", "scheduler", "temperature", "gamma", "beta1", "beta2"}
|
|
RESULT_HINT = re.compile(r"(?i)result|metric|score|eval|summary|log|history|progress|benchmark|runs?|ablation|output")
|
|
|
|
|
|
def walk(root):
|
|
for p in sorted(root.rglob("*")):
|
|
rel = p.relative_to(root).parts
|
|
if any(part in SKIP or (part.startswith(".") and part != ".") for part in rel[:-1]):
|
|
continue
|
|
if p.is_file():
|
|
yield p
|
|
|
|
|
|
def read(p, limit=2_000_000):
|
|
try:
|
|
return p.read_text(encoding="utf-8", errors="replace")[:limit]
|
|
except OSError:
|
|
return ""
|
|
|
|
|
|
def py_info(src, name):
|
|
info = {"imports": set(), "entry_point": False, "cli_args": [], "seeds": [], "functions": [], "classes": []}
|
|
try:
|
|
tree = ast.parse(src)
|
|
except SyntaxError:
|
|
tree = None
|
|
if tree:
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Import):
|
|
info["imports"] |= {a.name.split(".")[0] for a in node.names}
|
|
elif isinstance(node, ast.ImportFrom) and node.module and node.level == 0:
|
|
info["imports"].add(node.module.split(".")[0])
|
|
elif isinstance(node, ast.FunctionDef) and node.col_offset == 0:
|
|
info["functions"].append(node.name)
|
|
elif isinstance(node, ast.ClassDef) and node.col_offset == 0:
|
|
bases = [ast.unparse(b) for b in node.bases]
|
|
info["classes"].append(f"{node.name}({', '.join(bases)})" if bases else node.name)
|
|
elif isinstance(node, ast.Call) and getattr(node.func, "attr", "") == "add_argument":
|
|
flags = [a.value for a in node.args if isinstance(a, ast.Constant) and isinstance(a.value, str)]
|
|
kw = {k.arg: ast.unparse(k.value) for k in node.keywords if k.arg in ("default", "type", "help", "choices")}
|
|
if flags:
|
|
info["cli_args"].append({"flag": flags[-1], **{k: v[:80] for k, v in kw.items()}})
|
|
info["entry_point"] = bool(re.search(r"__name__\s*==\s*['\"]__main__['\"]|@hydra\.main|@click\.command|typer\.run", src))
|
|
info["seeds"] = sorted(set(SEED_RE.findall(src)))[:10]
|
|
return info
|
|
|
|
|
|
def nb_source(text):
|
|
try:
|
|
nb = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
return "", []
|
|
code, outs = [], []
|
|
for c in nb.get("cells", []):
|
|
s = "".join(c.get("source", []))
|
|
if c.get("cell_type") == "code":
|
|
code.append(s)
|
|
for o in c.get("outputs", []):
|
|
t = "".join(o.get("text", []) or o.get("data", {}).get("text/plain", []))
|
|
if t and LOG_METRIC_RE.search(t):
|
|
outs.append(t.strip()[:300])
|
|
elif c.get("cell_type") == "markdown":
|
|
code.append("# " + s.replace("\n", "\n# "))
|
|
return "\n".join(code), outs[:10]
|
|
|
|
|
|
def table_info(p, text):
|
|
if p.suffix in (".csv", ".tsv"):
|
|
lines = [l for l in text.splitlines() if l.strip()]
|
|
if not lines:
|
|
return None
|
|
sep = "\t" if p.suffix == ".tsv" or lines[0].count("\t") > lines[0].count(",") else ","
|
|
return {"kind": "table", "columns": [c.strip() for c in lines[0].split(sep)][:40], "rows": len(lines) - 1,
|
|
"head": lines[1:4]}
|
|
if p.suffix == ".jsonl":
|
|
lines = text.splitlines()
|
|
try:
|
|
first = json.loads(lines[0])
|
|
except (json.JSONDecodeError, IndexError):
|
|
return None
|
|
return {"kind": "jsonl", "keys": list(first)[:40] if isinstance(first, dict) else None, "rows": len(lines),
|
|
"head": lines[:2]}
|
|
if p.suffix == ".json":
|
|
try:
|
|
js = json.loads(text)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
def shape(x, d=0):
|
|
if d > 2: return "..."
|
|
if isinstance(x, dict): return {k: shape(v, d + 1) for k, v in list(x.items())[:15]}
|
|
if isinstance(x, list): return [f"list[{len(x)}]", shape(x[0], d + 1) if x else None]
|
|
return type(x).__name__ if not isinstance(x, (int, float)) else x
|
|
return {"kind": "json", "shape": shape(js)}
|
|
return None
|
|
|
|
|
|
def is_config(text, suffix):
|
|
keys = set(re.findall(r"(?m)^\s*[\"']?([A-Za-z_][\w.-]*)[\"']?\s*[:=]", text))
|
|
return len({k.lower().split(".")[-1] for k in keys} & HP_KEYS) >= 2
|
|
|
|
|
|
def scan(root):
|
|
inv = {"root": str(root.resolve()), "code": [], "notebooks": [], "configs": [], "results": [], "logs": [],
|
|
"tensorboard": [], "trackers": [], "checkpoints": Counter(), "docs": [], "figures": [],
|
|
"libraries": Counter(), "datasets": Counter(), "models": Counter(), "optimizers": Counter(),
|
|
"metrics": Counter()}
|
|
for d in ("wandb", "mlruns", "lightning_logs", "runs", "outputs", "multirun"):
|
|
for p in root.rglob(d):
|
|
if p.is_dir() and not (set(p.relative_to(root).parts) & SKIP):
|
|
inv["trackers"].append(str(p.relative_to(root)))
|
|
for p in walk(root):
|
|
rel, suf, name = str(p.relative_to(root)), p.suffix.lower(), p.name.lower()
|
|
if suf in (".pt", ".pth", ".ckpt", ".safetensors", ".h5", ".keras", ".pkl", ".joblib", ".onnx", ".bin", ".npz", ".npy"):
|
|
inv["checkpoints"][suf] += 1; continue
|
|
if "tfevents" in name:
|
|
inv["tensorboard"].append(rel); continue
|
|
if suf in (".png", ".jpg", ".jpeg", ".pdf", ".svg"):
|
|
inv["figures"].append(rel); continue
|
|
if name.startswith("readme") or name in ("requirements.txt", "environment.yml", "pyproject.toml", "setup.py", "agents.md") or suf == ".md":
|
|
text = read(p, 20_000)
|
|
inv["docs"].append({"path": rel, "first_lines": [l for l in text.splitlines() if l.strip()][:6]})
|
|
inv["datasets"].update(m.group(1) for m in DATASET_RE.finditer(text))
|
|
inv["models"].update(MODEL_RE.findall(text))
|
|
continue
|
|
if suf in (".py", ".ipynb"):
|
|
text = read(p)
|
|
outs = []
|
|
if suf == ".ipynb":
|
|
text, outs = nb_source(text)
|
|
info = py_info(text, rel)
|
|
libs = sorted(i for i in info["imports"] if i in ML_LIBS)
|
|
inv["libraries"].update(ML_LIBS[l] for l in libs)
|
|
ds = Counter(m.group(1) for m in DATASET_RE.finditer(text))
|
|
ds.update(":".join(x for x in m.groups() if x) for m in HF_DS_RE.finditer(text))
|
|
models = Counter(MODEL_RE.findall(text))
|
|
optims = Counter(OPTIM_RE.findall(text))
|
|
metrics = Counter(m.lower() for m in METRIC_RE.findall(text))
|
|
inv["datasets"].update(ds); inv["models"].update(models); inv["optimizers"].update(optims)
|
|
if not (set(libs) & CORE_ML) and not ds and not info["entry_point"]:
|
|
continue # plain helper without ML content
|
|
inv["metrics"].update(metrics)
|
|
entry = {"path": rel, "lines": text.count("\n") + 1, "ml_libraries": [ML_LIBS[l] for l in libs],
|
|
"entry_point": info["entry_point"], "cli_args": info["cli_args"][:40], "seeds": info["seeds"],
|
|
"datasets": sorted(ds), "models": sorted(models), "optimizers": sorted(optims),
|
|
"metrics": sorted(metrics), "functions": info["functions"][:30], "classes": info["classes"][:20]}
|
|
if outs:
|
|
entry["metric_outputs"] = outs
|
|
inv["notebooks" if suf == ".ipynb" else "code"].append(entry)
|
|
continue
|
|
if suf in (".yaml", ".yml", ".toml", ".cfg", ".ini") or (suf == ".json" and not RESULT_HINT.search(rel)):
|
|
text = read(p, 200_000)
|
|
if is_config(text, suf):
|
|
inv["configs"].append({"path": rel, "preview": [l for l in text.splitlines() if l.strip()][:25]})
|
|
inv["datasets"].update(m.group(1) for m in DATASET_RE.finditer(text))
|
|
continue
|
|
if suf != ".json":
|
|
continue
|
|
if suf in (".csv", ".tsv", ".json", ".jsonl"):
|
|
text = read(p)
|
|
ti = table_info(p, text)
|
|
if ti:
|
|
cols = " ".join(ti.get("columns") or ti.get("keys") or [str(ti.get("shape"))[:500]])
|
|
inv["metrics"].update(m.lower() for m in METRIC_RE.findall(cols))
|
|
inv["results"].append({"path": rel, "bytes": p.stat().st_size, **ti})
|
|
continue
|
|
if suf in (".log", ".txt", ".out", ".err"):
|
|
text = read(p)
|
|
hits = LOG_METRIC_RE.findall(text)
|
|
if hits:
|
|
inv["logs"].append({"path": rel, "metric_lines": len(hits),
|
|
"metrics": sorted({h[0].lower() for h in hits})[:20],
|
|
"last": [l for l in text.splitlines() if LOG_METRIC_RE.search(l)][-3:]})
|
|
for k in ("libraries", "datasets", "models", "optimizers", "metrics", "checkpoints"):
|
|
inv[k] = dict(inv[k].most_common(40))
|
|
return inv
|
|
|
|
|
|
def to_md(inv):
|
|
L = [f"# Experiment inventory\n\nRoot: `{inv['root']}`\n"]
|
|
for k in ("libraries", "datasets", "models", "optimizers", "metrics"):
|
|
L.append(f"- **{k}**: " + (", ".join(f"{a} ({n})" for a, n in inv[k].items()) or "-"))
|
|
L.append(f"- **checkpoints**: {inv['checkpoints'] or '-'} | **tensorboard files**: {len(inv['tensorboard'])}"
|
|
f" | **trackers**: {', '.join(inv['trackers']) or '-'}\n")
|
|
L.append("## Code (entry points first)\n")
|
|
for c in sorted(inv["code"] + inv["notebooks"], key=lambda c: (not c["entry_point"], c["path"])):
|
|
L.append(f"### `{c['path']}` ({c['lines']} lines{', ENTRY POINT' if c['entry_point'] else ''})")
|
|
L.append(f"libs: {', '.join(c['ml_libraries']) or '-'}; datasets: {', '.join(c['datasets']) or '-'}; "
|
|
f"models: {', '.join(c['models']) or '-'}; optimizers: {', '.join(c['optimizers']) or '-'}; "
|
|
f"metrics: {', '.join(c['metrics']) or '-'}; seeds: {', '.join(c['seeds']) or '-'}")
|
|
if c["cli_args"]:
|
|
L.append("args: " + "; ".join(f"{a['flag']}={a.get('default', '?')}" for a in c["cli_args"]))
|
|
if c["classes"]:
|
|
L.append("classes: " + ", ".join(c["classes"]))
|
|
if c["functions"]:
|
|
L.append("functions: " + ", ".join(c["functions"]))
|
|
for o in c.get("metric_outputs", []):
|
|
L.append(f" output> {o[:200]}")
|
|
L.append("")
|
|
L.append("## Configs\n")
|
|
for c in inv["configs"]:
|
|
L.append(f"### `{c['path']}`\n```\n" + "\n".join(c["preview"]) + "\n```")
|
|
L.append("\n## Result files\n")
|
|
for r in inv["results"]:
|
|
desc = (f"columns={r['columns']}, rows={r['rows']}" if r["kind"] == "table" else
|
|
f"keys={r['keys']}, rows={r['rows']}" if r["kind"] == "jsonl" else f"shape={json.dumps(r['shape'])[:400]}")
|
|
L.append(f"- `{r['path']}` ({r['kind']}, {r['bytes']} B): {desc}")
|
|
L.append("\n## Logs with metrics\n")
|
|
for g in inv["logs"]:
|
|
L.append(f"- `{g['path']}`: {g['metric_lines']} metric lines, {g['metrics']}; last: {g['last'][-1][:150] if g['last'] else ''}")
|
|
if inv["tensorboard"]:
|
|
L.append("\n## TensorBoard event files\n" + "\n".join(f"- `{t}`" for t in inv["tensorboard"][:50]))
|
|
if inv["figures"]:
|
|
L.append("\n## Existing figures (not reused blindly - regenerate from data)\n" + "\n".join(f"- `{f}`" for f in inv["figures"][:50]))
|
|
L.append("\n## Docs\n")
|
|
for d in inv["docs"]:
|
|
L.append(f"- `{d['path']}`: " + " / ".join(d["first_lines"])[:300])
|
|
return "\n".join(L) + "\n"
|
|
|
|
|
|
def main(argv):
|
|
if "-h" in argv or "--help" in argv:
|
|
print(__doc__); return 0
|
|
out = Path(argv[argv.index("--out") + 1]) if "--out" in argv else Path("paper/inventory")
|
|
pos = [a for i, a in enumerate(argv) if not a.startswith("--") and (i == 0 or argv[i - 1] != "--out")]
|
|
root = Path(pos[0] if pos else ".")
|
|
inv = scan(root)
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.with_suffix(".json").write_text(json.dumps(inv, indent=1, ensure_ascii=False, default=list), encoding="utf-8")
|
|
out.with_suffix(".md").write_text(to_md(inv), encoding="utf-8")
|
|
entries = [c["path"] for c in inv["code"] + inv["notebooks"] if c["entry_point"]]
|
|
print(f"wrote {out}.md and {out}.json")
|
|
print(f"code files with ML content: {len(inv['code'])} (+{len(inv['notebooks'])} notebooks); entry points: {', '.join(entries[:12]) or '-'}")
|
|
print(f"configs: {len(inv['configs'])}; result files: {len(inv['results'])}; logs with metrics: {len(inv['logs'])}; "
|
|
f"tensorboard: {len(inv['tensorboard'])}; trackers: {', '.join(inv['trackers']) or '-'}")
|
|
for k in ("libraries", "datasets", "models", "metrics"):
|
|
print(f"{k}: " + (", ".join(list(inv[k])[:12]) or "-"))
|
|
if not inv["results"] and not inv["logs"] and not inv["tensorboard"]:
|
|
print("NOTE: no result files found - the experiments may not have been run yet.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv[1:]))
|