Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Offline inspection and controlled GEPA optimization using Space-compatible LM secrets. | |
| Requires DSPy only for optimize. All run output stays under an explicit run directory. | |
| Sample JSONL rows: {"question": "Classify ...", "answer": "complaint"}. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| from datetime import datetime, timezone | |
| import importlib.metadata | |
| import json | |
| import logging | |
| import os | |
| from pathlib import Path | |
| import sys | |
| import threading | |
| import time | |
| from typing import Any | |
| def timestamp() -> str: | |
| return datetime.now(timezone.utc).isoformat() | |
| def atomic_json(path: Path, payload: dict[str, Any]) -> None: | |
| temp = path.with_name(path.name + ".tmp") | |
| temp.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8") | |
| os.replace(temp, path) | |
| def read_json(path: Path) -> dict[str, Any]: | |
| if not path.exists(): | |
| return {} | |
| return json.loads(path.read_text(encoding="utf-8")) | |
| def update_status(run_dir: Path, **fields: Any) -> dict[str, Any]: | |
| status_file = run_dir / "status.json" | |
| status = read_json(status_file) | |
| status.update(fields) | |
| status["updated_at"] = timestamp() | |
| atomic_json(status_file, status) | |
| return status | |
| def normalize(value: Any) -> str: | |
| return " ".join(str(value or "").casefold().split()) | |
| def load_rows(path: Path) -> list[dict[str, str]]: | |
| rows: list[dict[str, str]] = [] | |
| if not path.is_file(): | |
| raise ValueError(f"Dataset not found: {path}") | |
| for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): | |
| if not line.strip(): | |
| continue | |
| item = json.loads(line) | |
| if not isinstance(item, dict) or not isinstance(item.get("question"), str) or not isinstance(item.get("answer"), str): | |
| raise ValueError(f"{path}:{lineno}: expected string question and answer") | |
| if not item["question"].strip(): | |
| raise ValueError(f"{path}:{lineno}: question must not be empty") | |
| rows.append({"question": item["question"], "answer": item["answer"]}) | |
| if not rows: | |
| raise ValueError(f"Dataset is empty: {path}") | |
| return rows | |
| def assert_disjoint(train: list[dict[str, str]], val: list[dict[str, str]], test: list[dict[str, str]]) -> None: | |
| sets = [set(normalize(row["question"]) for row in rows) for rows in (train, val, test)] | |
| if sets[0] & sets[1] or sets[0] & sets[2] or sets[1] & sets[2]: | |
| raise ValueError("Dataset leakage: questions overlap train, val, or holdout test") | |
| def score_answer(expected: str, actual: str, metric: str) -> float: | |
| exp, got = normalize(expected), normalize(actual) | |
| if metric == "exact": | |
| return float(exp == got) | |
| if metric == "contains": | |
| return float(bool(exp) and exp in got) | |
| raise ValueError(f"Unsupported metric: {metric}") | |
| def dspy_model_name(name: str) -> str: | |
| # COMPATIBLE_MODEL is the endpoint's native model ID, not a DSPy provider ID. | |
| return name if name.startswith("openai/") else "openai/" + name | |
| def preflight() -> int: | |
| keys = ("COMPATIBLE_API", "COMPATIBLE_URL", "COMPATIBLE_MODEL") | |
| try: | |
| version = importlib.metadata.version("dspy") | |
| except importlib.metadata.PackageNotFoundError: | |
| version = None | |
| result = { | |
| "status": "ready" if version and all(os.environ.get(k) for k in keys) else "blocked", | |
| "dspy_version": version, | |
| "python": sys.version.split()[0], | |
| "secrets_present": {key: bool(os.environ.get(key)) for key in keys}, | |
| "reflection_model_configured": bool(os.environ.get("DSPY_REFLECTION_MODEL")), | |
| "default_threads": 1, | |
| "note": "This preflight is offline and does not verify provider inference.", | |
| } | |
| print(json.dumps(result, indent=2)) | |
| return 0 if result["status"] == "ready" else 2 | |
| def inspect(run_dir: Path) -> int: | |
| status = read_json(run_dir / "status.json") | |
| if not status: | |
| print(json.dumps({"status": "missing", "run_dir": str(run_dir)})) | |
| return 2 | |
| paths = ["optimized.json", "events.log", "gepa_logs"] | |
| status["artifacts"] = {p: (run_dir / p).exists() for p in paths} | |
| print(json.dumps(status, indent=2, sort_keys=True)) | |
| return 0 if status.get("state") == "completed" else 2 | |
| def evaluate(program: Any, examples: list[Any], metric_name: str) -> dict[str, Any]: | |
| correct = 0.0 | |
| for example in examples: | |
| output = program(question=example.question) | |
| correct += score_answer(example.answer, getattr(output, "answer", ""), metric_name) | |
| return {"score": round(correct / len(examples), 6), "examples": len(examples)} | |
| class MetricMonitor: | |
| """GEPA feedback metric; store aggregate progress, not example text or credentials.""" | |
| def __init__(self, dspy: Any, run_dir: Path, metric_name: str): | |
| self.dspy = dspy | |
| self.run_dir = run_dir | |
| self.metric_name = metric_name | |
| self.count = 0 | |
| self.lock = threading.Lock() | |
| def __call__(self, example: Any, pred: Any, trace: Any = None) -> Any: | |
| actual = str(getattr(pred, "answer", "") or "") | |
| expected = str(example.answer) | |
| score = score_answer(expected, actual, self.metric_name) | |
| feedback = ( | |
| "Matches the expected task output." | |
| if score == 1.0 else | |
| "Does not satisfy the reference output. Expected: " + expected[:200] + | |
| "; received: " + actual[:200] + | |
| ". Improve task instructions and consistent output format." | |
| ) | |
| with self.lock: | |
| self.count += 1 | |
| update_status(self.run_dir, metric_calls=self.count, last_score=score, phase="gepa") | |
| return self.dspy.Prediction(score=score, feedback=feedback) | |
| def _make_lm(dspy: Any, model: str, url: str, key: str, *, reflection: bool) -> Any: | |
| return dspy.LM( | |
| dspy_model_name(model), | |
| api_base=url, | |
| api_key=key, | |
| temperature=0.7 if reflection else 0.0, | |
| max_tokens=2048 if reflection else 512, | |
| ) | |
| def optimize(args: argparse.Namespace) -> int: | |
| # Validate dataset and environment before making any provider requests. | |
| train_rows = load_rows(args.train) | |
| val_rows = load_rows(args.val) | |
| test_rows = load_rows(args.test) if args.test else [] | |
| assert_disjoint(train_rows, val_rows, test_rows) | |
| values = {k: os.environ.get(k, "").strip() for k in ("COMPATIBLE_API", "COMPATIBLE_URL", "COMPATIBLE_MODEL")} | |
| missing = [key for key, value in values.items() if not value] | |
| if missing: | |
| raise RuntimeError("Missing required Space environment variables: " + ", ".join(missing)) | |
| if args.num_threads < 1 or args.num_threads > 8: | |
| raise ValueError("num-threads must be within 1..8") | |
| if args.max_metric_calls is not None and args.max_metric_calls < 1: | |
| raise ValueError("max-metric-calls must be positive") | |
| if args.run_dir.exists(): | |
| raise FileExistsError(f"Run directory already exists: {args.run_dir}. Use a new run ID.") | |
| old_umask = os.umask(0o077) | |
| try: | |
| args.run_dir.mkdir(parents=True, exist_ok=False) | |
| finally: | |
| os.umask(old_umask) | |
| started = time.monotonic() | |
| update_status( | |
| args.run_dir, state="running", phase="initializing", started_at=timestamp(), | |
| pid=os.getpid(), metric=args.metric, train_examples=len(train_rows), | |
| validation_examples=len(val_rows), test_examples=len(test_rows), | |
| metric_calls=0, budget=args.auto if args.max_metric_calls is None else args.max_metric_calls, | |
| num_threads=args.num_threads, | |
| ) | |
| try: | |
| import dspy | |
| logger = logging.getLogger("dspy") | |
| logger.setLevel(logging.INFO) | |
| file_handler = logging.FileHandler(args.run_dir / "events.log", encoding="utf-8") | |
| logger.addHandler(file_handler) | |
| try: | |
| student_lm = _make_lm(dspy, values["COMPATIBLE_MODEL"], values["COMPATIBLE_URL"], values["COMPATIBLE_API"], reflection=False) | |
| reflection_lm = _make_lm( | |
| dspy, | |
| os.environ.get("DSPY_REFLECTION_MODEL") or values["COMPATIBLE_MODEL"], | |
| os.environ.get("DSPY_REFLECTION_URL") or values["COMPATIBLE_URL"], | |
| os.environ.get("DSPY_REFLECTION_API") or values["COMPATIBLE_API"], | |
| reflection=True, | |
| ) | |
| dspy.configure(lm=student_lm) | |
| trainset = [dspy.Example(**row).with_inputs("question") for row in train_rows] | |
| valset = [dspy.Example(**row).with_inputs("question") for row in val_rows] | |
| testset = [dspy.Example(**row).with_inputs("question") for row in test_rows] | |
| baseline = dspy.Predict("question -> answer") | |
| update_status(args.run_dir, phase="baseline") | |
| baseline_validation = evaluate(baseline, valset, args.metric) | |
| update_status(args.run_dir, phase="gepa", baseline_validation=baseline_validation) | |
| monitor = MetricMonitor(dspy, args.run_dir, args.metric) | |
| kwargs: dict[str, Any] = { | |
| "metric": monitor, | |
| "reflection_lm": reflection_lm, | |
| "log_dir": str(args.run_dir / "gepa_logs"), | |
| "num_threads": args.num_threads, | |
| "track_stats": True, | |
| "track_best_outputs": True, | |
| "use_merge": False, | |
| } | |
| if args.max_metric_calls is None: | |
| kwargs["auto"] = args.auto | |
| else: | |
| kwargs["max_metric_calls"] = args.max_metric_calls | |
| optimizer = dspy.GEPA(**kwargs) | |
| optimized = optimizer.compile(student=baseline, trainset=trainset, valset=valset) | |
| optimized.save(str(args.run_dir / "optimized.json")) | |
| update_status(args.run_dir, phase="validation") | |
| optimized_validation = evaluate(optimized, valset, args.metric) | |
| holdout = None | |
| if testset: | |
| update_status(args.run_dir, phase="holdout") | |
| holdout = { | |
| "baseline": evaluate(baseline, testset, args.metric), | |
| "optimized": evaluate(optimized, testset, args.metric), | |
| } | |
| result = update_status( | |
| args.run_dir, state="completed", phase="done", finished_at=timestamp(), | |
| elapsed_seconds=round(time.monotonic() - started, 3), | |
| optimized_validation=optimized_validation, | |
| validation_delta=round(optimized_validation["score"] - baseline_validation["score"], 6), | |
| holdout=holdout, | |
| holdout_verified=bool(testset), | |
| optimized_artifact="optimized.json", | |
| ) | |
| print(json.dumps(result, indent=2, sort_keys=True)) | |
| return 0 | |
| finally: | |
| logger.removeHandler(file_handler) | |
| file_handler.close() | |
| except Exception as exc: | |
| # Do not persist exception messages: some providers include credentials in errors. | |
| update_status( | |
| args.run_dir, state="failed", phase="error", finished_at=timestamp(), | |
| elapsed_seconds=round(time.monotonic() - started, 3), error_type=type(exc).__name__, | |
| ) | |
| raise | |
| def main() -> int: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| subs = parser.add_subparsers(dest="command", required=True) | |
| subs.add_parser("preflight", help="Offline dependency and secret presence checks") | |
| watch = subs.add_parser("inspect", help="Offline run status and artifact inspection") | |
| watch.add_argument("--run-dir", type=Path, required=True) | |
| run = subs.add_parser("optimize", help="Run a GEPA experiment") | |
| run.add_argument("--train", type=Path, required=True) | |
| run.add_argument("--val", type=Path, required=True) | |
| run.add_argument("--test", type=Path) | |
| run.add_argument("--run-dir", type=Path, required=True) | |
| run.add_argument("--metric", choices=("exact", "contains"), default="exact") | |
| budget = run.add_mutually_exclusive_group() | |
| budget.add_argument("--auto", choices=("light", "medium", "heavy"), default="light") | |
| budget.add_argument("--max-metric-calls", type=int) | |
| run.add_argument("--num-threads", type=int, default=1) | |
| args = parser.parse_args() | |
| if args.command == "preflight": | |
| return preflight() | |
| if args.command == "inspect": | |
| return inspect(args.run_dir) | |
| # Keep all generated artifacts and optimizer trace logs private even when | |
| # the parent shell has a permissive umask. | |
| previous_umask = os.umask(0o077) | |
| try: | |
| return optimize(args) | |
| except Exception as exc: | |
| # stderr should not accidentally contain provider-returned secrets. | |
| print("GEPA run failed: " + type(exc).__name__ + ". See the run status for evidence.", file=sys.stderr) | |
| return 1 | |
| finally: | |
| os.umask(previous_umask) | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |