File size: 4,045 Bytes
dc9e1b9 | 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 | """Build a flat run directory in the layout the experiments were run in.
Every script here loads lenses, prompts and the jacobian-lens experiment data
by bare relative path from the current directory (e.g.
"jacobian_lens_ouro-2.6b_t47.pt", "jacobian-lens/data/experiments/..."), and
writes its outputs there too. This repo stores lenses under lenses/<model>/,
so this script symlinks everything into one directory to run from.
The Ouro-2.6B-Thinking lenses share filenames with the base Ouro lenses, so
they need their own run directory (--thinking), with OURO_MODEL_ID set.
Lenses are symlinked (read-only inputs); results are copied, because the
scripts rewrite their outputs by bare filename and would otherwise write
through a link into this repo. To refit lenses, use --no-lenses: the merge
scripts save jacobian_lens_*.pt by bare filename too.
Usage:
python code/setup_workdir.py RUN_DIR --jlens PATH/TO/jacobian-lens
[--thinking] [--with-results] [--no-lenses]
--jlens a git checkout of https://github.com/anthropics/jacobian-lens
(at the commit in revisions.JLENS_COMMIT)
--thinking link the Ouro-2.6B-Thinking lenses instead of the base
Qwen/Ouro/Huginn ones
--with-results also copy in the released results, so analysis/plot
steps can run without regenerating them on a GPU
--no-lenses link no lenses (for refitting them from scratch)
"""
import argparse
import shutil
import subprocess
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "code"))
from revisions import JLENS_COMMIT # noqa: E402
def link(src, dst):
if dst.is_symlink() or dst.exists():
if dst.is_symlink() and dst.resolve() == src.resolve():
return
raise SystemExit(f"refusing to overwrite {dst}")
dst.symlink_to(src.resolve())
def copy(src, dst):
if dst.is_symlink() or dst.exists():
raise SystemExit(f"refusing to overwrite {dst}")
if src.is_dir():
shutil.copytree(src, dst)
else:
shutil.copy2(src, dst)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("run_dir")
ap.add_argument("--jlens", required=True)
ap.add_argument("--thinking", action="store_true")
ap.add_argument("--with-results", action="store_true")
ap.add_argument("--no-lenses", action="store_true")
args = ap.parse_args()
run = Path(args.run_dir)
run.mkdir(parents=True, exist_ok=True)
jlens = Path(args.jlens)
if not (jlens / "data" / "experiments").is_dir():
raise SystemExit(f"{jlens} is not a jacobian-lens checkout "
"(no data/experiments/)")
try:
head = subprocess.run(["git", "-C", str(jlens), "rev-parse", "HEAD"],
capture_output=True, text=True, check=True).stdout.strip()
except (OSError, subprocess.CalledProcessError):
head = None
if head != JLENS_COMMIT:
print(f"WARNING: jacobian-lens is at {head}, the runs used {JLENS_COMMIT}")
link(jlens, run / "jacobian-lens")
for f in sorted((ROOT / "code").iterdir()):
if f.suffix in (".py", ".sh") and f.name != "setup_workdir.py":
link(f, run / f.name)
models = ["ouro_thinking"] if args.thinking else ["qwen", "ouro", "huginn"]
if args.no_lenses:
models = []
for m in models:
for f in sorted((ROOT / "lenses" / m).glob("*.pt")):
link(f, run / f.name)
if args.with_results:
res = ROOT / "results" / ("thinking" if args.thinking else "")
for f in sorted(res.iterdir()):
if f.name in ("thinking", "logs"):
continue
copy(f, run / f.name)
print(f"run directory ready: {run}")
if args.thinking:
print("before running, export OURO_MODEL_ID=ByteDance/Ouro-2.6B-Thinking")
print("generate the fitting prompts there first: python prep_prompts.py")
if __name__ == "__main__":
main()
|