import logging
import argparse
import json
import time
from pathlib import Path
from subprocess import run, PIPE  # type: ignore
from typing import Dict, Any, List, Optional

def run_cmd(cmd: List[str]) -> Dict[str, Any]:
	start = time.time()
	p = run(cmd, stdout=PIPE, stderr=PIPE, encoding="utf-8")
	dur = time.time() - start
	ok = False
	risk = None
	mem_hit = None
	try:
		lines = p.stdout.strip().splitlines() if p.stdout else []
		# Look for the line with "status" field (not just the last line)
		obj = {}
		for line in lines:


import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--domain", choices=["cli", "dom"], default="cli")
	ap.add_argument("--wm", default=str(root / "models" / "wm_mlp" / "wm_mlp.joblib"))
	ap.add_argument("--il", default=str(root / "data" / "skills" / "actuator_il.jsonl"))
	ap.add_argument("--t5", default=str(root / "models" / "actuator_il_t5"))
	ap.add_argument("--obs", required=True)
	ap.add_argument("--plan", required=True)
	args = ap.parse_args()

	obs = json.loads(args.obs)
	plan = json.loads(args.plan)

	from agi_dw.core.world_model.api import WorldModelPrior  # type: ignore


import logging
import argparse
import json

from agi_dw.bench.web_dom.runner import fetch_text

def main() -> None:
	parser = argparse.ArgumentParser()
	parser.add_argument("--url", required=True)
	parser.add_argument("--selector", default="body")
	args = parser.parse_args()
	res = fetch_text(args.url, args.selector)
	print(json.dumps(res, ensure_ascii=False))

if __name__ == "__main__":
	main()


import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--skill-id", required=True)
	ap.add_argument("--verifier-adapter", default=None)
	ap.add_argument("--planner-adapter", default=None)
	args = ap.parse_args()

	from agi_dw.core.memory.skills import SkillLibrary  # type: ignore

	lib = SkillLibrary(str(root))
	adapters = {}
	if args.verifier_adapter:
		adapters["verifier"] = str(args.verifier_adapter)
	if args.planner_adapter:


import logging
import argparse
import json
import time
from pathlib import Path
from subprocess import run, PIPE  # type: ignore

def _verify_cmd(root: Path, src: Path, dst: Path, backend: str, batch_size: int) -> list[str]:
	cmd = [
		"python3", str(root / "scripts" / "verify_traces.py"),
		str(src),
		str(dst),
		"--backend", backend,
		"--require-llm",
		"--structured", "json",
		"--warmup",
		"--workers", "1",
		"--batch-size", str(int(batch_size)),
		"--max", "20",
	]


import logging
import argparse
import json
from pathlib import Path
from typing import List, Dict, Any

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--include-cli", action="store_true")
	ap.add_argument("--include-dom", action="store_true")
	ap.add_argument("--out", default=str(root / "data" / "benchmarks" / "external_suite.jsonl"))
	args = ap.parse_args()

	recs: List[Dict[str, Any]] = []
	if args.include_cli:
		recs.append({"domain": "cli", "task": {"name": "count_lines", "timeout": 20}})
		recs.append({"domain": "cli", "task": {"name": "grep_error", "timeout": 20}})
	if args.include_dom:
		# Pull from YAML seeds if available


#!/usr/bin/env python3
import logging
import argparse
import json
import re
from pathlib import Path
from typing import List, Dict, Any

def parse_args():
	root = Path(__file__).resolve().parents[1]
	ap = argparse.ArgumentParser(description="Simple contamination checker for code samples")
	ap.add_argument("--samples", nargs="*", default=[
		str(root / "data" / "llm_bench" / "humaneval_samples.jsonl"),
		str(root / "data" / "llm_bench" / "mbpp_samples.jsonl"),
		str(root / "data" / "llm_bench" / "apps_samples.jsonl"),
		str(root / "data" / "llm_bench" / "ds1000_samples.jsonl"),
	])
	ap.add_argument("--corpus", default=str(root / "data" / "corpora" / "train_corpus.txt"), help="Plain text file to scan against")
	ap.add_argument("--out", default=str(root / "data" / "llm_bench" / "contamination.json"))
	return ap.parse_args()


import math, copy, random
from dataclasses import dataclass, field
from typing import Dict, Any, List, Tuple, Iterable, Callable

import torch
import torch.nn as nn
import torch.nn.functional as F

# ============================================================
# 0) Demo model + toy data (replace with your LLM & loaders)
# ============================================================

class TinyMLP(nn.Module):
    def __init__(self, d=2, h=64, out=2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d, h), nn.GELU(),
            nn.Linear(h, h), nn.GELU(),
            nn.Linear(h, out),
        )


import logging
import argparse
import json
from pathlib import Path
from typing import Dict, Any, List

def predict_action(obs: Dict[str, Any], plan: Dict[str, Any], dom_t5: str, dom_il: str) -> Dict[str, Any]:
	try:
		from agi_dw.core.actuator.t5_actuator import ActuatorT5Predictor  # type: ignore
		act = ActuatorT5Predictor(dom_t5, mode="dom", structured=True)
		a = act.predict_action(obs, plan)
		if isinstance(a, dict):
			return a
	except Exception:
		pass
	try:
		from agi_dw.core.actuator.il_baseline import ActuatorILNearestNeighbor  # type: ignore
		il = ActuatorILNearestNeighbor(dom_il)
		a = il.predict_action(obs, plan)
		if isinstance(a, dict):


import logging
import argparse
import json
import re
from pathlib import Path
from typing import Dict, List, Set, Tuple

LEGACY_TO_NAMESPACED_MAP = {
	"ci-probes": "ci.probes",
	"ci-matrix": "ci.matrix",
	"ci-weekly": "ci.weekly",
	"code-index": "tools.code-index",
	"index-scripts": "tools.index-scripts",
	"index-docs": "tools.index-docs",
	"run-loop-oscli": "loops.oscli",
	"run-loop-webdom": "loops.webdom",
	"loop-oscli": "loops.oscli",
	"loop-webdom": "loops.webdom",
	"build-balanced-splits": "splits.balanced",
	"build-planner-splits": "splits.planner",


#!/usr/bin/env python3
import logging
import argparse
import json
from pathlib import Path
from subprocess import run, PIPE

def parse_args():
	root = Path(__file__).resolve().parents[1]
	ap = argparse.ArgumentParser(description="Run lm-eval-harness (if installed) and save JSON metrics")
	ap.add_argument("--model", default="meta-llama/Llama-3.2-3B")
	ap.add_argument("--tasks", default="mmlu,hellaswag,winogrande,piqa,arc_challenge,boolq")
	ap.add_argument("--out", default=str(root / "data" / "llm_bench" / "lmeval.json"))
	ap.add_argument("--limit", default=None, help="Optional eval limit per task (e.g., 100)")
	return ap.parse_args()

def main() -> int:
	args = parse_args()
	out = Path(args.out)


import logging
import argparse
import json
from pathlib import Path
from datetime import datetime

from agi_dw.bench.web_dom.runner import fetch_text  # type: ignore
from agi_dw.bench.common.trace import build_trace, write_jsonl  # type: ignore
from agi_dw.core.verifier.llm_verifier import verify_trace_snippet  # type: ignore

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--url", default="https://example.com")
	ap.add_argument("--selector", default="h1")
	ap.add_argument("--out", default=str(root / "data" / "traces" / "web_dom.jsonl"))
	ap.add_argument("--verifier-backend", choices=["hf"], default="hf")
	ap.add_argument("--verifier-model", default="meta-llama/Llama-3.2-3B")
	ap.add_argument("--structured", choices=["none", "json"], default="json")
	ap.add_argument("--adapter", default=None)


#!/usr/bin/env python3
"""
Massive-scale task extractor for the inspiration folder.
This script extracts training data from all major repositories and converts them
into our AGI training format. This is a large-scale operation that could generate
thousands of training examples.
"""

import json
import os
import random
import re
from pathlib import Path
from typing import Dict, List, Any, Optional
import logging

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MassiveTaskExtractor:


import math, random, copy, os
from dataclasses import dataclass, field
from typing import List, Dict, Callable

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_

# ---------------------------
# 0) Tiny dataset (char-level)
# ---------------------------
def make_char_data(text: str, block_size=64, device="cpu"):
    vocab = sorted(list(set(text)))
    stoi = {ch: i for i, ch in enumerate(vocab)}
    itos = {i: ch for ch, i in stoi.items()}

    def encode(s):
        return torch.tensor([stoi[c] for c in s], dtype=torch.long)


import logging
import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def generate_plan_candidates(obs: Dict[str, Any], model: str, num_candidates: int = 3) -> List[Dict[str, Any]]:
	from agi_dw.core.llm.hf_client import HFClient  # type: ignore
	client = HFClient.get_cached(model)
	prompt = (
		f"You are a planner. Given the observation below, propose {num_candidates} diverse candidate plans as a JSON array.\n"
		"Each plan must be an object with keys: subgoals (array of strings), tools (array of strings), constraints (object).\n"
		"Return ONLY JSON (no code fences).\n\nObservation (JSON):\n" + json.dumps(obs, ensure_ascii=False)
	)
	text = client.generate(prompt, max_new_tokens=400, temperature=0.2)
	try:
		arr = json.loads(text)
		if isinstance(arr, list):
			out: List[Dict[str, Any]] = []
			for it in arr:


import logging
import argparse
import json
import sys
from pathlib import Path
from typing import Any, Dict, List

def ensure_curriculum(root: Path, path: Path) -> None:
	if not path.exists():
		try:
			from scripts.misc.generate_practice_curriculum import main as gen_main  # type: ignore
		except Exception:
			return
		try:
			gen_main()
		except SystemExit:
			pass

def load_tier(path: Path, tier: str) -> List[Dict[str, Any]]:


#!/usr/bin/env python3
import logging
import argparse
import json
import os
import tempfile
from pathlib import Path

def parse_args():
	root = Path(__file__).resolve().parents[1]
	ap = argparse.ArgumentParser(description="Run style/type checks on multi-language code samples (if tools installed)")
	ap.add_argument("--samples", nargs="*", default=[
		str(root / "data" / "llm_bench" / "humaneval_samples.jsonl"),
		str(root / "data" / "llm_bench" / "mbpp_samples.jsonl"),
		str(root / "data" / "llm_bench" / "apps_samples.jsonl"),
		str(root / "data" / "llm_bench" / "ds1000_samples.jsonl"),
		str(root / "data" / "llm_bench" / "cruxeval_samples.jsonl"),
	])
	ap.add_argument("--out", default=str(root / "data" / "llm_bench" / "code_style_multi.json"))
	return ap.parse_args()


import logging
import argparse
import json
import random
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Any, List, Optional, Tuple

import numpy as np
from tqdm import tqdm

from agi_dw.core.world_model.rollout_api import RolloutAPI, RolloutConfig  # type: ignore
from agi_dw.core.world_model.validator_reward import ValidatorRewardShaper, RewardConfig  # type: ignore

@dataclass
class TrainingConfig:
	"""Configuration for off-policy training."""
	batch_size: int = 32
	num_epochs: int = 10


import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--path", default=str(root / "models" / "memory"))
	ap.add_argument("--query", required=True)
	ap.add_argument("--topk", type=int, default=5)
	ap.add_argument("--recency", type=float, default=0.0)
	ap.add_argument("--quality", type=float, default=0.0, help="Quality weighting [0..1] to favor high-quality items")
	ap.add_argument("--max-age-days", type=float, default=None, help="Discard items older than this many days")
	ap.add_argument("--include-tags", nargs='*', default=[])
	ap.add_argument("--exclude-tags", nargs='*', default=[])
	args = ap.parse_args()

	from agi_dw.core.memory.episodic import EpisodicMemory  # type: ignore


import logging
import argparse
import json
from pathlib import Path
from typing import Dict, Iterable

def iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				yield json.loads(line)
			except Exception:
				continue


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def merge_traces(inputs: List[Path], out_path: Path, max_items: int | None = None) -> int:
	out_path.parent.mkdir(parents=True, exist_ok=True)
	n = 0
	with out_path.open("w", encoding="utf-8") as out:
		for p in inputs:
			try:
				with p.open("r", encoding="utf-8") as f:
					for line in f:
						if not line.strip():
							continue
						out.write(line)
						n += 1


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def _safe(obj: Any) -> Any:
	try:
		json.dumps(obj, ensure_ascii=False)
		return obj
	except Exception:
		return str(obj)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--in", dest="inp", default=str(root / "data" / "devtools" / "traces.jsonl"))


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, Iterable, Tuple

def iter_jsonl(path: Path) -> Iterable[Dict[str, Any]]:
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			s = line.strip()
			if not s:
				continue
			try:
				yield json.loads(s)
			except Exception:
				continue


import logging
import argparse
import hashlib
import json
from pathlib import Path
from typing import Dict, Tuple

def iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				yield json.loads(line)
			except Exception:
				continue


import logging
import argparse
import json
from pathlib import Path
from typing import Dict, Any, Tuple

def iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				yield json.loads(line)
			except Exception:
				continue


import logging
import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def _extract_effects(trace: Dict[str, Any]) -> str:
	# Prefer domain-specific effects fields
	res = trace.get("result", {}) or {}
	if isinstance(res, dict):
		# DOM text
		dom_text = res.get("dom")
		if isinstance(dom_text, str) and dom_text:
			return dom_text[:1000]
		# CLI stdout
		stdout = res.get("stdout")
		if isinstance(stdout, str) and stdout:
			return stdout[:1000]
	# Fallback: empty
	return ""


from __future__ import annotations
import logging

import argparse
from pathlib import Path
import sys

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--inputs", nargs="*", default=[
		str(root / "data" / "traces" / "dev_loop.jsonl"),
		str(root / "data" / "traces" / "dev_loop.ci.jsonl"),
	])
	ap.add_argument("--out", default=str(root / "data" / "traces" / "nearmiss.jsonl"))
	args = ap.parse_args()

	# Delegate to existing near-miss replay builder
	cmd = [
		sys.executable,


import logging
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List
import sys

def build_example(obj: Dict[str, Any]) -> Dict[str, Any]:
	obs = obj.get("obs", {})
	plan = obj.get("plan", {})
	action = obj.get("action", {})
	input_text = json.dumps({"obs": obs, "plan": plan}, ensure_ascii=False)
	output_text = json.dumps(action, ensure_ascii=False)
	return {"input": input_text, "output": output_text}

def iter_jsonl(paths: Iterable[Path]) -> Iterable[Dict[str, Any]]:
	for p in paths:
		if not p.exists():
			continue
		for line in p.read_text(encoding="utf-8").splitlines():


import logging
import argparse
import json
from pathlib import Path
from typing import List

def load_traces(paths: List[str]) -> List[dict]:
	out = []
	for p in paths:
		fp = Path(p)
		if not fp.exists():
			continue
		with fp.open("r", encoding="utf-8") as f:
			for line in f:
				line = line.strip()
				if not line:
					continue
				try:
					out.append(json.loads(line))
				except Exception:


import logging
import json
from pathlib import Path
from typing import Any, Dict
import sys

def build_example(obj: Dict[str, Any]) -> Dict[str, Any]:
	obs = obj.get("obs", {})
	plan = obj.get("plan", {})
	action = obj.get("action", {})
	# Minimal text input for IL: obs + plan as JSON
	input_text = json.dumps({"obs": obs, "plan": plan}, ensure_ascii=False)
	output_text = json.dumps(action, ensure_ascii=False)
	return {"input": input_text, "output": output_text}

def process(src_jsonl: Path, out_jsonl: Path) -> int:
	out_jsonl.parent.mkdir(parents=True, exist_ok=True)
	count = 0
	with src_jsonl.open("r", encoding="utf-8") as fin, out_jsonl.open("w", encoding="utf-8") as fout:


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict

def _safe(obj: Any) -> Any:
	try:
		json.dumps(obj, ensure_ascii=False)
		return obj
	except Exception:
		return str(obj)

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--in", dest="inp", required=True, help="Path to EBST JSONL input")


import logging
import argparse
import hashlib
import json
from pathlib import Path
from typing import Dict, Iterable, Tuple

def iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				yield json.loads(line)
			except Exception:
				continue


import logging
import argparse
import json
from pathlib import Path
from typing import Dict, Any

from agi_dw.core.verifier.service import VerifierServiceConfig, verify as verifier_run  # type: ignore

from agi_dw.core.actuator.service import compute_router_features  # type: ignore

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument(
		"--verified",
		nargs="+",
		default=[
			str(root / "data" / "traces" / "seed_os_cli.verified.llm.jsonl"),
		],
		help="One or more verified JSONL files (CLI and/or DOM)",


import logging
import argparse
import json
from pathlib import Path
from typing import Dict

def build_example(trace: Dict) -> Dict[str, str]:
	# Compact input: obs+plan+result only
	inp = {
		"obs": trace.get("obs", {}),
		"plan": trace.get("plan", {}),
		"result": trace.get("result", {}),
	}
	# Heuristic target YAML (bootstrap)
	status = (trace.get("result") or {}).get("status", "error")
	if status == "ok":
		success_prob, risk = 0.8, 0.2
	else:
		success_prob, risk = 0.3, 0.6
	tgt = (


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Dict, Any

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--out", default=str(root / "data" / "traces" / "practice_ds.jsonl"))
	ap.add_argument("--n", type=int, default=5, help="Number of synthetic practice items")
	ap.add_argument("--seeds", default="1,2,3,4,5", help="Comma-separated integer seeds to vary sampling")
	args = ap.parse_args()

	outp = Path(args.out)
	outp.parent.mkdir(parents=True, exist_ok=True)

	# Minimal synthetic practice items: trivial failing tests with gold diff


import logging
import argparse
import json
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple

def iter_jsonl(path: Path) -> Iterable[Dict[str, Any]]:
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			line = line.strip()
			if not line:
				continue
			try:
				yield json.loads(line)
			except Exception:
				continue


import logging
import argparse
import json
from pathlib import Path
from typing import List

def discover_verified_traces(data_root: Path) -> List[str]:
	cands: List[str] = []
	for name in ["seed_os_cli.verified.jsonl", "seed_os_cli.verified.llm.jsonl", "web_dom.verified.jsonl"]:
		p = data_root / "traces" / name
		if p.exists():
			cands.append(str(p))
	return cands

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--data-root", default=str(root / "data"))
	ap.add_argument("--out", default=str(root / "models" / "memory"))


import logging
import argparse
import json
from pathlib import Path
from typing import Any, Dict

from agi_dw.bench.web_dom.runner import fetch_text
from agi_dw.bench.common.trace import build_trace, write_jsonl
from agi_dw.core.planner.service import PlannerConfig, VerifierConfig, WMConfig, ContextAugment, plan_with_context
from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RouterVerifierConfig, WMPriorConfig, WMScreenConfig, RepairConfig, select_action
from agi_dw.core.verifier.llm_verifier import verify_trace_snippet
from agi_dw.core.memory.service import match_skill_action
from datetime import datetime
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore

def main() -> int:
	parser = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	parser.add_argument("--url", default="https://www.iana.org/domains/reserved")


import logging
import argparse
import json
from pathlib import Path
from datetime import datetime
import os
from typing import Dict
import subprocess
import shlex

# Make script robust if PYTHONPATH is not set: add project root to sys.path
try:
	from agi_dw.tools.git import GitTool
	from agi_dw.tools.test_runner import TestRunner
	from agi_dw.core.actuator.patch_actuator import PatchActuator
	from agi_dw.tools.failure_classifier import classify_failures
	from agi_dw.tools.linter import LinterTool
	from agi_dw.core.llm.hf_client import HFClient
	from agi_dw.core.utils.prompt_logger import PromptLogger, get_prompt_logger  # type: ignore
	from agi_dw.core.actuator.code_actions import execute_code_action  # type: ignore
	from agi_dw.bench.common.trace import build_trace, write_jsonl  # type: ignore


import logging
import argparse
import json
from pathlib import Path
from typing import Any, Dict, Tuple, Optional, List

from agi_dw.bench.common.safe_shell import SafeShellRunner
from datetime import datetime
from agi_dw.bench.common.trace import build_trace, write_jsonl
from agi_dw.bench.os_cli.tasks import setup_count_lines, setup_grep_word
from agi_dw.core.planner.service import PlannerConfig, VerifierConfig, WMConfig, ContextAugment, plan_with_context
from agi_dw.core.ops.tracing import trace_span, meter_cost  # type: ignore
from agi_dw.core.actuator.service import ActuatorConfig, RouterExtras, RouterVerifierConfig, WMPriorConfig, WMScreenConfig, RepairConfig, select_action
from agi_dw.core.memory.service import match_skill_action

def prepare_task(sandbox: Path, task: str) -> Tuple[Dict[str, Any], str]:
	if task == "count_lines":
		setup_count_lines(sandbox)
		obs = {"kind": "cli", "content": "Count file lines", "meta": {"cwd": str(sandbox)}}
		return obs, "count_lines"


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def iter_jsonl(path: Path):
	if not path.exists():
		return
	with path.open("r", encoding="utf-8") as f:
		for line in f:
			s = line.strip()
			if not s:
				continue
			try:
				yield json.loads(s)
			except Exception:
				continue


import logging
import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

def _safe_float(x: Any, default: float = 0.0) -> float:
	try:
		return float(x)
	except Exception:
		return float(default)

def _p90(values: List[float]) -> float:
	if not values:
		return 0.0
	vals = sorted(values)
	idx = int(min(len(vals) - 1, max(0, round(0.9 * (len(vals) - 1)))))
	return float(vals[idx])


from __future__ import annotations
import logging

import argparse
import json
from pathlib import Path

def _read_json(path: Path) -> dict:
	try:
		if path.exists():
			return json.loads(path.read_text(encoding="utf-8"))
	except Exception:
		pass
	return {}

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[2]
	ap.add_argument("--devloop", default=str(root / "data" / "dashboards" / "devloop.json"))


import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--traces", default=str(root / "data" / "traces" / "loop_run.jsonl"))
	ap.add_argument("--out", default=str(root / "data" / "dashboards" / "trajectory_summary.json"))
	args = ap.parse_args()

	trp = Path(args.traces)
	if not trp.exists():
		print(json.dumps({"ok": False, "error": "traces_missing", "path": str(trp)}))
		return 1

	n_total = 0
	n_success = 0
	steps_total = 0


#!/usr/bin/env python3
import logging
import argparse
import json
from pathlib import Path

def parse_args():
	root = Path(__file__).resolve().parents[1]
	ap = argparse.ArgumentParser(description="Aggregate LLM benchmark results into dashboard section")
	ap.add_argument("--results", default=str(root / "data" / "llm_bench" / "results.json"))
	ap.add_argument("--dashboard", default=str(root / "data" / "dashboards" / "summary.json"))
	ap.add_argument("--section", default="llm_bench")
	ap.add_argument("--lmeval", default=str(root / "data" / "llm_bench" / "lmeval.json"))
	return ap.parse_args()

def main() -> int:
	args = parse_args()
	results_path = Path(args.results)
	dash_path = Path(args.dashboard)


import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	ap.add_argument("--verified", required=True)
	ap.add_argument("--out", default=None)
	args = ap.parse_args()

	path = Path(args.verified)
	if not path.exists():
		print("verified file not found:", str(path))
		return 2
	total = 0
	ok = 0
	latencies = []
	risks = []
	with path.open("r", encoding="utf-8") as f:


#!/usr/bin/env python3
import logging
import argparse
import json
from pathlib import Path

def parse_args():
	root = Path(__file__).resolve().parents[1]
	ap = argparse.ArgumentParser(description="Aggregate coding benchmark results into dashboard section")
	ap.add_argument("--results", default=str(root / "data" / "llm_bench" / "code_results.json"))
	ap.add_argument("--dashboard", default=str(root / "data" / "dashboards" / "summary.json"))
	ap.add_argument("--section", default="code_bench")
	ap.add_argument("--style", default=str(root / "data" / "llm_bench" / "code_style.json"))
	return ap.parse_args()

def main() -> int:
	args = parse_args()
	dash_path = Path(args.dashboard)
	dash_path.parent.mkdir(parents=True, exist_ok=True)


import logging
import argparse
import json
from collections import defaultdict
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--verified", default=str(root / "data" / "traces" / "web_dom.verified.jsonl"))
	ap.add_argument("--out", default=str(root / "data" / "traces" / "web_dom.summary.json"))
	args = ap.parse_args()

	path = Path(args.verified)
	outp = Path(args.out)
	if not path.exists():
		print(json.dumps({"ok": False, "error": "verified_missing", "path": str(path)}))
		return 1

	total = 0


import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--bench", default=str(root / "data" / "benchmarks" / "kpi.json"))
	ap.add_argument("--wm", default=str(root / "models" / "wm_mlp" / "metrics.json"))
	ap.add_argument("--verifier", default=str(root / "models" / "verifier_calib" / "metrics.json"))
	ap.add_argument("--devloop", default=str(root / "data" / "dashboards" / "devloop.json"))
	ap.add_argument("--coder-rl", default=str(root / "models" / "coder_rl" / "metrics.json"))
	ap.add_argument("--coder-nearmiss", default=str(root / "models" / "coder_nearmiss" / "model.json"))
	ap.add_argument("--planner", default=str(root / "data" / "planner_prefs" / "metrics.json"))
	ap.add_argument("--offpolicy", default=str(root / "models" / "wm_offpolicy" / "metrics.json"))
	ap.add_argument("--shaping", default=str(root / "data" / "planner" / "reward_shaping.json"))
	ap.add_argument("--out", default=str(root / "data" / "dashboards" / "summary.json"))
	ap.add_argument("--use-budgeted", action="store_true", help="Prefer budgeted success rates if available")
	ap.add_argument("--traces", default=str(root / "data" / "traces" / "loop_run.jsonl"), help="Optional traces JSONL to compute planner self-eval summary")


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict, List

def callers_of(index_graph: Dict[str, Any], symbol: str) -> List[str]:
    res: List[str] = []
    calls = (index_graph.get("calls") or {}) if isinstance(index_graph, dict) else {}
    for f, call_list in calls.items():
        try:
            if any((c.get("name") == symbol) for c in (call_list or [])):
                res.append(f)
        except Exception:
            continue
    return res


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Dict, Any, List

def shard_by_top_level(root: Path, index_graph: Dict[str, Any]) -> Dict[str, List[str]]:
    shards: Dict[str, List[str]] = {}
    files = list((index_graph.get("functions") or {}).keys()) if isinstance(index_graph, dict) else []
    for f in files:
        try:
            p = Path(f)
            top = (p.parts or [""])[0]
        except Exception:
            top = ""
        shards.setdefault(top, []).append(f)
    return shards


#!/usr/bin/env python3
from __future__ import annotations
import logging

import ast
import sys
from pathlib import Path

TEMPLATE = """
from agi_dw.flags.sdk import FLAGS as _FLAGS

def {name}_dual_route(*args, **kwargs):
    if _FLAGS.is_enabled('{flag_name}', default=True):
        return {primary}(*args, **kwargs)
    return {secondary}(*args, **kwargs)
"""

def inject_dual_route(path: Path, symbol_primary: str, symbol_secondary: str, flag_name: str) -> bool:
    src = path.read_text(encoding="utf-8")


#!/usr/bin/env python3
from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Any, List

def inject_logging(text: str) -> str:
	if "import logging" in text:
		return text
	# Insert at top after shebang/encoding if present
	lines = text.splitlines()
	insert_at = 0
	if lines and lines[0].startswith("#!/"):
		insert_at = 1
	lines.insert(insert_at, "import logging")
	return "\n".join(lines)

def main() -> int:


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, Any, List, Tuple

def run_gates(repo_dir: Path) -> Tuple[bool, Dict[str, Any]]:
    """Run a minimal gate suite: flake/mypy/pytest if present; summarize exit codes."""
    summary: Dict[str, Any] = {}
    ok = True
    cmds = [
        ("pytest -q", "tests"),
        ("flake8 .", "."),
        ("mypy .", "."),
    ]


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict

def _load_json(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		return json.loads(path.read_text(encoding="utf-8"))
	except Exception:
		return {}

def _safe_load_yaml(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	plan_path = root / "release" / "plan.yaml"
	status_path = root / "data" / "sandbox" / "tmp" / f"release_{datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')}.json"
	# Simulate a publish according to plan
	plan_text = plan_path.read_text(encoding="utf-8") if plan_path.exists() else "{}"
	status: Dict[str, Any] = {
		"ok": True,
		"plan": str(plan_path),
		"timestamp": datetime.utcnow().isoformat() + "Z",
		"result": {"canary": "started", "notes": "simulated publish"},


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	status_path = root / "data" / "sandbox" / "tmp" / f"release_rollback_{datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')}.json"
	status: Dict[str, Any] = {
		"ok": True,
		"timestamp": datetime.utcnow().isoformat() + "Z",
		"result": {"rollback": "completed", "notes": "simulated rollback"},
	}
	status_path.parent.mkdir(parents=True, exist_ok=True)
	status_path.write_text(json.dumps(status, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(status_path)}))


from typing import List, Dict, Any, Tuple
import sys, io, re, json, math, statistics, itertools, functools, copy, random
import torch
from collections import Counter
from agi_dw.core.llm.hf_client import HFClient  # type: ignore
from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
from transformers.generation.logits_process import (  # type: ignore
    LogitsProcessor,
    LogitsProcessorList,
)
from agi_dw.core.llm.adapter_cache import AdapterCache  # type: ignore

import os
import subprocess
import shlex
from pathlib import Path
from agi_dw.scripts.selfplay.modules.paths import get_data_root, data_path  # type: ignore
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, TimeoutError
from agi_dw.scripts.selfplay.modules.lora import _loraify_linear, _inject_lora, _forward_with_lora_delta, _set_lora_enabled, _patch_forwards


import os
import types
import torch

from agi_dw.scripts.selfplay.modules import tasks, generation, sandbox, healing

class FakeTokenizer:
    def __init__(self, texts):
        self._texts = list(texts)
        self.eos_token_id = 0
        self.pad_token_id = 0

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):


from .common_imports import *
from .text_utils import extract_first_fenced_block
import ast

def _safe_builtins() -> Dict[str, Any]:
    # Build a minimal allowed builtins dict and expose it both as __builtins__ and globals
    allowed: Dict[str, Any] = {}
    for k in [
        "abs","all","any","enumerate","len","range","min","max","sum","map","filter",
        # add numeric and sequence helpers commonly needed
        "zip","sorted","reversed","list","tuple","dict","set", "bool", "int", "float", "str", "print",
        # allow stdlib imports (episode gate handles policy)
        "__import__"
    ]:
        try:
            allowed[k] = getattr(__builtins__, k)
        except Exception:
            pass
    b: Dict[str, Any] = {"__builtins__": allowed}
    b.update(allowed)
    return b


from __future__ import annotations

from pathlib import Path
from typing import Any, Optional
import os

def _find_repo_root(anchor: Optional[Path] = None) -> Path:
    p = (anchor or Path(__file__)).resolve()
    for parent in [p] + list(p.parents):
        try:
            if parent.name == "agi_dw":
                return parent
        except Exception:
            continue
    # Fallback heuristics: go up 3 levels from modules/, else 2
    try:
        return p.parents[3]
    except Exception:
        return p.parents[2]


from .common_imports import *
from .text_utils import extract_first_fenced_block

def _trace_io(kind: str, data: Dict[str, Any]) -> None:
    try:
        root = Path(__file__).resolve().parents[3]
        p = root / "scripts" / "data" / "traces" / ("dev_loop.jsonl" if kind == "text" else "loop_run.jsonl")
        p.parent.mkdir(parents=True, exist_ok=True)
        obj = {"kind": kind, **data}
        with p.open("a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")
    except Exception:
        pass

class SanitizeLogits(LogitsProcessor):
    """Clamp and sanitize logits to avoid NaN/Inf during sampling on CUDA.

    Replaces NaN with a large negative and +/-Inf with finite limits, then clamp tightly
    """

    def __call__(self, input_ids: torch.LongTensor, scores: torch.Tensor) -> torch.Tensor:  # type: ignore[override]


from .common_imports import *

def strip_non_code_preamble(src: str) -> str:
    try:
        s = (src or "").lstrip()
        if s.lower().startswith("python"):
            s = s[len("python"):].lstrip()
        s = s.replace("```python", "").replace("```", "")
        return s.strip()
    except Exception:
        return src

def ensure_top_import(src: str, line: str) -> str:
    try:
        lines = [l for l in (src or "").splitlines() if not l.strip().startswith("#")]
        if any(l.strip() == line for l in lines):
            return src
        for i, l in enumerate(lines):
            if l.strip().startswith("def solve("):


from .common_imports import *

def extract_first_fenced_block(text: str) -> str | None:
    """Return the contents of the first fenced code block, sans optional language header.

    Heuristic:
    - Find the first opening fence ```
    - Slice from the first newline after the opening to the first subsequent closing fence
    - Drop a one-line language header like 'python'/'python3'/'py' if present
    - Return stripped text or None if not found
    """
    try:
        s = str(text).lstrip("\ufeff")
        i0 = s.find("```")
        if i0 == -1:
            return None
        first_nl = s.find("\n", i0)
        if first_nl == -1:
            return None
        i1 = s.find("```", first_nl + 1)
        if i1 == -1:


from .common_imports import *
from dataclasses import dataclass

# Lightweight staged verifier with partial credit shaping

@dataclass
class StageScore:
    format_ok: float
    compile_ok: float
    signature_ok: float
    style_ok: float
    tests_pass: float

_FENCE_RE = re.compile(r"^```", re.MULTILINE)

def looks_like_pure_code(src: str) -> bool:
    try:


from .common_imports import *
from .generation import generate_text
from .train import sft_step

def mutate_prompt(prompt: str) -> str:
    muts = [
        " Optimize for speed.",
        " Make it memory efficient.",
        " Add error handling.",
        " Make it object-oriented.",
        " Include unit tests.",
    ]
    try:
        return (prompt or "").strip() + random.choice(muts)
    except Exception:
        return prompt

def heuristic_reward(output: str, prompt: str, max_gen_len: int) -> float:
    try:
        score = 0.0
        keywords = ["def", "class", "return", ":", "(", ")"]


from .common_imports import *

def _dev_episode(repo: str, extra_args: str | None = None, timeout_sec: int = 1200) -> Tuple[bool, Dict[str, Any]]:
    """Run one dev loop episode on a repo; return (ok, info)."""
    root = Path(__file__).resolve().parents[3]
    dev_script = root / 'scripts' / 'loops' / 'run_loop_dev.py'
    if not dev_script.exists():
        return False, {"error": "dev_script_not_found", "path": str(dev_script)}
    cmd = f"python {dev_script} --repo {repo} --apply-meta"
    if extra_args and extra_args.strip():
        cmd = f"{cmd} {extra_args.strip()}"
    try:
        p = subprocess.run(shlex.split(cmd), cwd=str(root), capture_output=True, text=True, timeout=timeout_sec)
        ok = (p.returncode == 0)
        info = {"returncode": p.returncode, "stdout": p.stdout[-2000:], "stderr": p.stderr[-1000:]}
        return ok, info
    except Exception as e:
        return False, {"error": str(e)}

def _ensure_local_repo(local_uri: str) -> str:
    """Ensure a local:/abs/path repo directory exists; return normalized URI."""


from .common_imports import *
try:
    from .verifier_stages import StageScore  # type: ignore
except Exception:
    StageScore = None  # type: ignore

def _compute_reward(meta: Dict[str, Any]) -> Tuple[float, Dict[str, float]]:
    """Compute composite reward from rep/verifier/WM/latency/policy signals.

    Returns (reward, components) with reward in [0,1].
    """
    rep = meta.get("rep", {}) if isinstance(meta, dict) else {}
    passed = float(rep.get("passed", 0) or 0)
    total = float(rep.get("total", 0) or 0)
    pass_ratio = (passed / total) if total > 0 else 0.0
    wm_risk = float(meta.get("wm_risk", meta.get("v_risk", 0.5) or 0.5))
    verifier_sp = float(meta.get("v_sp", meta.get("success_prob", 0.5) or 0.5))
    timeouts = float(rep.get("timeouts", 0) or 0)
    timeout_pen = (timeouts / total) if total > 0 else 0.0
    try:
        timeout_sec = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "1.5") or 1.5)


from .common_imports import *

class Memory:
    def __init__(self, path: str) -> None:
        self.path = path

    def log(self, obj: Dict[str, Any]) -> None:
        try:
            p = Path(self.path)
            p.parent.mkdir(parents=True, exist_ok=True)
        except Exception:
            pass
        with open(self.path, "a", encoding="utf-8") as f:
            f.write(json.dumps(obj, ensure_ascii=False) + "\n")

class Curriculum:
    def __init__(self) -> None:
        from collections import deque

        self.level = 0
        self.win_hist = deque(maxlen=100)


from __future__ import annotations

from typing import Callable, Dict, Any, List, Tuple, Optional
import random

Rand = random.Random

class TaskSpec(Dict[str, Any]):
    """Keeps same shape: {name, signature, tests}."""
    pass

class TaskTemplate:
    def __init__(
        self,
        name: str,
        signature: str,
        sampler: Callable[[Rand], Tuple[List[List[Any]], Callable[..., Any]]],
        tags: List[str],


from .common_imports import *

def props_rev_str(f) -> bool:
    try:
        import string as _s
        import random as _r
        for _ in range(64):
            s = "".join(_r.choice(_s.printable[:90]) for __ in range(_r.randint(0, 64)))
            r = f(s)
            if (r[::-1] != s) or (len(r) != len(s)):
                return False
            if f(r) != s:
                return False
        return True
    except Exception:
        return False

def tests_sum_to_n(f) -> bool:
    try:


from .common_imports import *
from .generation import generate_text

def heal_code(tok, model, prompt: str, failing_code: str, failure_info: Dict[str, Any], max_new_tokens: int = 256) -> str:
     """Construct a healing prompt that explains why the submission failed and
    reminds the model of exact signature/output constraints.

    We include:
    - First exception string (if any)
    - First input/output mismatch example (if any)
    - Aggregate stats (passed/total, timeouts)
    - A restatement of the exact required signature parsed from the task prompt
    """
    # Extract concise failure signal
    try:
        err = failure_info.get("first_failure") or failure_info.get("err") or ""
    except Exception:
        err = ""

    # Aggregate stats if available
    stats = ""


from .common_imports import *
from .lora import _inject_lora, _patch_forwards, _set_lora_enabled

def sft_step(tok, model, prompt: str, target: str, lr: float, steps: int, mem=None) -> float | None:
    # Ensure LoRA adapters exist, are enabled, and only A/B are trainable; refresh optimizer
    def _has_lora(m: torch.nn.Module) -> bool:
        try:
            for _n, _mm in m.named_modules():
                if isinstance(_mm, torch.nn.Linear) and hasattr(_mm, "A") and hasattr(_mm, "B"):
                    return True
        except Exception:
            pass
        return False
    if not _has_lora(model):
        try:
            r = int(os.environ.get("SELFPLAY_ADAPTER_R", "8") or 8)
        except Exception:
            r = 8
        try:
            alpha = int(os.environ.get("SELFPLAY_ADAPTER_ALPHA", "16") or 16)
        except Exception:


from .common_imports import *

def _resolve_best_wm_model_path() -> str:
    try:
        root = Path(__file__).resolve().parents[1]
        cands = [
            root / "models" / "wm_mlp" / "best.joblib",
            root / "models" / "wm_mlp" / "latest.joblib",
            root / "models" / "wm_mlp" / "wm_mlp.joblib",
        ]
        for p in cands:
            if p.exists():
                return str(p)
    except Exception:
        pass
    return ""


from typing import List, Dict, Any, Tuple, Optional
import sys, io, re, json, math, statistics, itertools, functools, copy, random
import torch
from collections import Counter
from agi_dw.core.llm.hf_client import HFClient  # type: ignore
from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
from transformers.generation.logits_process import (  # type: ignore
    LogitsProcessor,
    LogitsProcessorList,
)
from agi_dw.core.llm.adapter_cache import AdapterCache  # type: ignore
import os
import subprocess
import shlex
from pathlib import Path
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, TimeoutError


# Package for selfplay modularized components


from .common_imports import *
from .text_utils import extract_first_fenced_block

from .healing import _strip_triple_quoted
from .model_io import load_seed, save_adapters, load_adapters
from .wm import _resolve_best_wm_model_path
from .curriculum import Memory, Curriculum, _maybe_gate_curriculum
from .evolution import evolution_loop, sft_step
from .tasks import sample_task, render_prompt
from .generation import sample, generate_text, beam_candidates, speculative_accept_reject
from .early_exit import generate_with_early_exit, EarlyExitHead
from .taskgen import model_generate_task
from .foundry_llm import llm_sample_task
from .reward import _compute_reward
from .tasks_bank import (
    task_bank_add,
    task_bank_pick,
    task_bank_stats,
    task_bank_update,
    task_hash,
    summarize_tests,


from __future__ import annotations

from .common_imports import *
from .generation import generate_text
from .tasks_foundry import _REGISTRY as _FOUNDRY_REG, _mk_tests as _FOUNDRY_MK, TaskTemplate as _FT  # type: ignore

def _list_templates(level: int, want_tags: list[str]) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []
    for t in _FOUNDRY_REG:
        try:
            name = getattr(t, "name", "")
            sig = getattr(t, "signature", "")
            tags = list(getattr(t, "tags", []) or [])
            diff = int(getattr(t, "difficulty", 0))
            ok_diff = diff <= min(5, int(level) + 2)
            ok_tags = (not want_tags) or any((tg in tags) for tg in want_tags)
            if ok_diff and ok_tags:
                items.append({"name": name, "signature": sig, "tags": tags, "difficulty": diff})
        except Exception:
            continue


from __future__ import annotations

from .common_imports import *
from .paths import data_path

def task_bank_path() -> Path:
    return data_path("selfplay", "task_bank.jsonl")

def parse_denylist() -> set[str]:
    try:
        dl = str(os.environ.get("SELFPLAY_TASK_DENYLIST", "")).strip()
        if not dl:
            return set()
        return set(x.strip() for x in dl.split(",") if x.strip())
    except Exception:
        return set()

def task_bank_load() -> Dict[str, Any]:


from .common_imports import *
from .lora import _inject_lora, _patch_forwards

def _build_device_map_and_max_memory() -> Tuple[Any, Optional[Dict[Any, str]]]:
    """Build (device_map, max_memory) using env overrides, preferring larger GPUs.

    Env controls:
    - HF_DEVICE_MAP: explicit device_map string (e.g., "auto", "balanced_low_0")
    - HF_MAX_MEMORY_JSON: JSON dict like {"0":"23GiB","1":"23GiB","2":"6GiB","cpu":"64GiB"}
    - SELFPLAY_AVOID_GPU_IDS: comma-separated GPU ids to de-prioritize; they get "1GiB"
    - SELFPLAY_MAX_MEM_HEADROOM_GIB: integer GiB to leave free per preferred GPU (default 1)
    """
    # device_map
    dm_env = os.environ.get("HF_DEVICE_MAP", None)
    device_map: Any = (dm_env if (dm_env is not None and str(dm_env).strip() != "") else ("auto" if torch.cuda.is_available() else None))
    # max_memory
    mm_env = os.environ.get("HF_MAX_MEMORY_JSON", None)
    if mm_env:
        try:
            mm = json.loads(mm_env)
            # normalize keys to int where possible


from .common_imports import *

class EarlyExitHead(torch.nn.Module):
    def __init__(self, hidden_size: int):
        super().__init__()
        self.classifier = torch.nn.Linear(hidden_size, 1)

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        # hidden_state: [B, H]
        return torch.sigmoid(self.classifier(hidden_state))

def generate_with_early_exit(
    tok,
    model,
    prompt: str,
    head: EarlyExitHead,
    *,
    layer_index: int = -1,
    threshold: float = 0.9,


from .common_imports import *
try:
    from .tasks_foundry import _REGISTRY as _FOUNDRY_REG, _mk_tests as _FOUNDRY_MK, TaskTemplate as _FT
except Exception:
    _FOUNDRY_REG = []  # type: ignore

_PRIMS = [
    {
        "name": "sum_to_n",
        "sig": "def solve(n:int)->int:",
        "gen": lambda: {"args": [__import__("random").randint(1, 5000)]},
        "ref": lambda n: n * (n + 1) // 2,
        "desc": "Return the sum of integers from 1 to n (inclusive).",
    },
    {
        "name": "is_prime",
        "sig": "def solve(n:int)->bool:",
        "gen": lambda: {"args": [__import__("random").randint(0, 5000)]},
        "ref": lambda n: n > 1 and all(n % d for d in range(2, int(n ** 0.5) + 1)),
        "desc": "Return True if n is prime, else False.",
    },


from __future__ import annotations

from .common_imports import *
from .generation import generate_text
from .healing import heal_code
from .sandbox import _worker_run_one_subproc

def _extract_code_block(text: str) -> str | None:
    s = text or ""
    if s.count("```") >= 2:
        try:
            parts = s.split("```")
            inner = max((seg for seg in parts[1::2]), key=lambda u: len(u), default="")
            if inner:
                lines = inner.splitlines()
                if lines and re.match(r"^[a-zA-Z0-9_+-]+$", lines[0].strip()):
                    inner = "\n".join(lines[1:])
                s = inner.strip()
        except Exception:
            s = s


from .common_imports import *

def _loraify_linear(module: torch.nn.Module, r: int = 8, alpha: int = 16) -> None:
    if not isinstance(module, torch.nn.Linear):
        return
    # Idempotent: if already has LoRA params, skip
    try:
        if hasattr(module, "A") and hasattr(module, "B"):
            return
    except Exception:
        pass
    w = module.weight
    in_f, out_f = w.shape[1], w.shape[0]
    # Create and register LoRA params/buffers on the same device/dtype as base weight
    A = torch.nn.Parameter(torch.zeros(r, in_f, device=w.device, dtype=w.dtype))
    B = torch.nn.Parameter(torch.zeros(out_f, r, device=w.device, dtype=w.dtype))
    module.register_parameter("A", A)
    module.register_parameter("B", B)
    # Register scaling as a buffer to ensure device moves with the module
    module.register_buffer("lora_scaling", torch.tensor(float(alpha / r), dtype=w.dtype, device=w.device))
    module.weight.requires_grad_(False)


#!/usr/bin/env python3
from __future__ import annotations

import argparse
import collections as C
import json
import os
import re
from typing import Dict, Any

def pct(a: int, b: int) -> float:
    return (float(a) / float(max(1, b))) * 100.0

def normalize_first_failure(text: str, max_len: int = 200) -> str:
    try:
        s = re.sub(r"\s+", " ", str(text)).strip()
        return s[:max_len]
    except Exception:
        return str(text)[:max_len]


import os
import sys
import json
from pathlib import Path

# Ensure project root is on sys.path so 'agi_dw' package imports work when run directly
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests, run_and_test
from agi_dw.scripts.selfplay.modules.healing import heal_code

def _bool_env(name: str, default: bool) -> bool:


import os
import sys
import json
import itertools
from pathlib import Path

# Ensure project root on sys.path
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests

def main() -> None:


from __future__ import annotations
import json
from pathlib import Path
import os
from typing import Dict, Any

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore

from agi_dw.scripts.selfplay.modules.paths import data_path
from agi_dw.scripts.selfplay.modules.train import dpo_step, kto_step

def _env_str(name: str, default: str) -> str:
    v = os.environ.get(name)
    return str(v) if v is not None and str(v).strip() != "" else str(default)

def _env_float(name: str, default: float) -> float:
    try:
        v = os.environ.get(name)


import os
import sys
import json
from pathlib import Path

# Ensure project root on sys.path
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.core.llm.hf_client import HFClient  # type: ignore
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample
from agi_dw.scripts.selfplay.modules.sandbox import run_subset_tests

def main() -> None:
    model_name = os.environ.get("HF_PROBE_MODEL", os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"))


import os
import sys
from pathlib import Path

# Ensure project root is on sys.path so 'agi_dw' package imports work when run directly
try:
    ROOT = Path(__file__).resolve().parents[4]
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
except Exception:
    pass

from agi_dw.scripts.selfplay.modules.model_io import load_seed
from agi_dw.scripts.selfplay.modules.tasks import sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import sample, generate_text

def main() -> None:
    cfg = {
        "model_name": os.environ.get("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "direct_model": False,