from __future__ import annotations

import sys
from pathlib import Path
import subprocess

EX_DIR = Path("/data/transformer_10/examples/program_conditioned_adapter")
SMOKE_PROG = Path("/data/transformer_10/examples/program_conditioned_adapter/smoke_repo")
ART_DIR = EX_DIR / "examples" / "python_repo_grounded_planning" / "artifacts" / "smoke_planning"

PG_BACKEND = "examples.scripts.python_repo_graph:PythonRepoGraph"
EMIT_MOD = "examples.python_repo_grounded_planning.emit_planning_knowledge"

def _run(argv: list[str]) -> int:
    print("[run]", " ".join(argv))
    return subprocess.call(argv)

def main() -> None:
    ART_DIR.mkdir(parents=True, exist_ok=True)


from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Dict, List, Tuple, Optional, Set
import importlib
import numpy as np

def _load_symbol(path: str):
    mod, _, attr = path.partition(":")
    m = importlib.import_module(mod)
    return getattr(m, attr)

def _write_jsonl(path: Path, rows: List[Dict]) -> None:
    with path.open("w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")


from typing import Any, Dict, List, Optional, Tuple
import os
import re
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
try:
    import pathspec  # type: ignore
except Exception:  # pragma: no cover
    pathspec = None  # type: ignore

from examples.scripts.code_graph import CodeGraph

def build_repo_embedding(
    repo_root: str,
    *,
    dim: int = EMBED_DIM_DEFAULT,
    seed: int = 0,
    include_text: bool = False,
    text_max_bytes: int = 0,
    max_text_tokens: int = 0,
    text_weight: float = 0.25,


from __future__ import annotations

from typing import List, Tuple

def load_mbpp_texts(max_n: int = 128, split: str = "train") -> Tuple[List[str], int]:
	"""
	Load MBPP via Hugging Face datasets and return a list of training texts.
	Each text is prompt + solution combined to form a code-aware LM sample.
	Returns (texts, total_available).
	"""
	try:
		from datasets import load_dataset  # type: ignore
	except Exception:
		return [], 0
	# Try common configs in order of availability
	ds = None
	for name in ("mbpp",):
		for subset in ("sanitized", None):
			try:
				if subset is None:


from __future__ import annotations

import os
import ast
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Iterable, Set

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str


from __future__ import annotations

from typing import Any, Dict, List

def enforce(policy: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
	"""
	Apply simple safety gates like 'apply_patch requires tests_green' etc.
	Returns {'ok': bool, 'missing': [..]}.
	"""
	reqs: List[str] = list(policy.get("requires", []))
	passed = []
	missing = []
	for r in reqs:
		if context.get(r, False):
			passed.append(r)
		else:
			missing.append(r)
	return {"ok": len(missing) == 0, "missing": missing}


from __future__ import annotations

import subprocess
from typing import Any, Dict, List, Optional

def run_benchmark_verifier(cmd: Optional[List[str]], env: Optional[Dict[str, str]] = None, timeout_sec: int = 120) -> bool:
	"""
	Runs an external verifier command (official checker) and returns True on success.
	When cmd is None, returns True (no-op) for smoke runs.
	"""
	if not cmd:
		return True
	try:
		rc = subprocess.call(cmd, env=env, timeout=timeout_sec)  # type: ignore
	except Exception:
		return False
	return rc == 0


from __future__ import annotations

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

def synth_benchmark_heads(adapters_dir: str, datasets: List[str]) -> List[str]:
	"""
	Smoke-level synthesis of benchmark-aware adapter heads.
	Creates tiny shard stubs under adapters/shards/* so run.py can pick them up.
	"""
	adir = Path(adapters_dir)
	shards_dir = adir / "shards"
	shards_dir.mkdir(parents=True, exist_ok=True)
	created: List[str] = []
	for spec in datasets:
		ds_name = spec.replace(":", "_").replace("/", "_")
		path = shards_dir / f"benchmark_{ds_name}_head.json"


from __future__ import annotations

import os
from typing import Iterable, List, Tuple, Dict, Optional, Set

from modules.program_graph import (
    Entity,
    Edge,
    Artifact,
    Span,
    ResolvedAnchor,
    EntityId,
)
from .codegraph_core import CodeGraph, CGEntity, CGEdge
from .repo_graph import RepoGraph, artifact_uri, program_id_for_repo, parse_program_uri

def _entity_uri(program_id: str, cg: CGEntity) -> str:
    if cg.kind == "module" or cg.kind == "test_module":
        resource = cg.name
        kind = "module"


from __future__ import annotations

import os
import re
import hashlib
from typing import Iterable, List, Tuple, Dict, Optional, Set

from modules.program_graph import (
    ProgramGraph,
    Entity,
    Edge,
    Artifact,
    Span,
    ResolvedAnchor,
    EntityId,
)

def program_id_for_repo(repo_root: str) -> str:
    base = os.path.basename(os.path.abspath(repo_root)) or "repo"
    return base


from __future__ import annotations

from typing import Any, Dict

def allocate(budget: Dict[str, Any], expected_gain: float) -> Dict[str, Any]:
	"""
	Greedily allocate tokens/time by expected verifier-gain per unit cost (placeholder).
	"""
	out = dict(budget)
	out["allocated"] = {"tokens": int(budget.get("tokens", 64000) * 0.5)}
	out["expected_gain"] = float(expected_gain)
	return out


from __future__ import annotations

import os
from typing import Iterable, List, Tuple, Dict, Optional

from modules.program_graph import (
	ProgramGraph,
	Entity,
	Edge,
	Artifact,
	Span,
	ResolvedAnchor,
	EntityId,
)
from .repo_graph import artifact_uri, program_id_for_repo, parse_program_uri

class DatasetProgramGraph(ProgramGraph):
	"""
	Minimal Dataset→ProgramGraph backend.
	Represents benchmark datasets (MBPP, SWE-bench, etc.) as entities:


from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

def _default_registry_path() -> str:
	artifacts_dir = examples_dir / "artifacts"
	artifacts_dir.mkdir(parents=True, exist_ok=True)
	return str(artifacts_dir / "adapter_registry.json")

def _read_json(path: str) -> Dict[str, Any]:
	if not os.path.isfile(path):
		return {}
	try:
		with open(path, "r", encoding="utf-8") as fh:
			return json.loads(fh.read()) or {}


from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Tuple, Iterable

# Canonical PCA Interface (portable utility helpers)

def select_subgraph(goal: Dict[str, Any], program_facts: Dict[str, Any]) -> Dict[str, Any]:
	"""
	Select(): question-aware selection of subgraph/segments/windows from the program.
	Heuristic: prefer recently changed files, failing spans, or symbols mentioned in goal text.
	"""
	text = (goal.get("text") or goal.get("prompt") or "").lower()
	return {"windows": [], "symbols": [], "goal_terms": text.split()[:16]}

def pack_with_anchors(sources: Iterable[Path], windows: List[Tuple[str, int, int]]) -> Dict[str, Any]:
	"""
	Pack(): deterministic context packaging with anchored snippets (path:line windows).


from __future__ import annotations

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

def _read_json(path: str) -> Dict[str, Any]:
	with open(path, "r", encoding="utf-8") as fh:
		return json.loads(fh.read())

def run_program_training(adapters_dir: str, plan_json: str, out_json: str) -> str:
	"""
	Smoke-level trainer hook used by examples: simulates a training pass
	grounded to a ProgramGraph-based TrainingPlan by producing a summary
	and a trained marker under adapters_dir.
	"""
	adapters_dir_abs = os.path.abspath(adapters_dir)


from __future__ import annotations

from typing import Any, Dict, List

def require_citations(outputs: Dict[str, Any], min_count: int = 1) -> Dict[str, Any]:
	"""
	Ensure outputs contain at least min_count citation anchors.
	"""
	cites = outputs.get("citations") or []
	ok = isinstance(cites, list) and len(cites) >= int(min_count)
	return {"ok": ok, "count": len(cites) if isinstance(cites, list) else 0}


import os
import ast
import re
import json
import time  # noqa: F401
from dataclasses import dataclass, field
from typing import Any, Dict, List, Tuple, Optional
try:
    import pathspec  # type: ignore
except Exception:  # pragma: no cover
    pathspec = None  # type: ignore

@dataclass
class Symbol:
    fqn: str
    name: str
    qualname: str
    kind: str  # module|class|function|variable
    module: str
    file: str


from __future__ import annotations

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

def _read_json(path: Path) -> Dict[str, Any]:
	try:
		with open(path, "r", encoding="utf-8") as fh:
			return json.loads(fh.read())
	except Exception:
		return {}

def _load_local_jsonl(fp: Path, text_key: str = "text", max_n: int | None = None) -> List[str]:
	texts: List[str] = []
	if not fp.exists():
		return texts
	with open(fp, "r", encoding="utf-8") as fh:


from __future__ import annotations

import os
import sys
import json
from pathlib import Path
from datetime import datetime, timezone
from typing import Dict, Any

def _resolve_pg_ctor(pg_backend: str):
    mod, _, attr = pg_backend.partition(":")
    if not mod or not attr:
        raise ValueError(f"Invalid ProgramGraph backend '{pg_backend}', expected 'module:ClassName'")
    m = __import__(mod, fromlist=[attr])
    return getattr(m, attr)

def emit_repository_knowledge(repo_root: str, out_path: str, pg_backend: str) -> str:
    """Emit a consolidated repository_knowledge.json with entities, edges, and artifact spans."""
    repo_root_abs = os.path.abspath(repo_root)


from __future__ import annotations

from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Tuple, Set, Any
import os
import json
import time
import math
import numpy as np

def _now_ts() -> float:
    try:
        return time.time()
    except Exception:
        return 0.0

def _unit(x: np.ndarray) -> np.ndarray:
    n = float(np.linalg.norm(x))
    return (x / n) if n > 0 else x


from __future__ import annotations

import os
import sys
import subprocess
from pathlib import Path

# Resolve program_conditioned_adapter directory directly from this file
EX_DIR = Path(__file__).resolve().parents[2]
SMOKE_REPO = EX_DIR / "smoke_repo"
# Backend class name (implemented in this example package)
PG_BACKEND = "examples.scripts.python_repo_graph:PythonRepoGraph"
# Program configuration loader
from .program_config import load_program_config  # noqa: E402

def _run(cmd: list[str]) -> int:
    print("[run]", " ".join(cmd))
    return subprocess.call(cmd)


from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

@dataclass(frozen=True)
class ProgramContracts:
    require_citations: bool = True
    citations_per_paragraph: bool = False
    retrieval_policy: str = "sim:0.6,struct:0.4"
    retrieval_temp: float = 0.7

@dataclass(frozen=True)
class ProgramPaths:
    adapters_dir: Path
    knowledge_path: Path
    program_state_path: Optional[Path] = None


from __future__ import annotations

import os
import ast
import hashlib
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Iterable, Set

@dataclass(frozen=True)
class FileSpan:
    file: str           # absolute path
    start_line: int     # 1-based inclusive
    end_line: int       # 1-based inclusive

@dataclass
class CGEntity:
    id: str             # stable id (e.g., fqn)
    kind: str           # module|function|class|test
    name: str


from typing import List
from ..code_graph import CodeGraph

def build_prompts_for_module(g: CodeGraph, module: str, max_q: int = 3) -> List[str]:
    """Construct a few simple, verifiable prompts for a module using symbol names and doc headers.

    Prefers questions that can be answered via local context/citations.
    """
    prompts: List[str] = []
    # Symbols defined in the module
    defs = list(g.defs_in(module) or [])
    # Ask for an explanation of the module
    prompts.append(f"Explain the key functions/classes in {module}. Cite path:line for each claim.")
    # Ask about up to two concrete defs
    for fqn in defs[:2]:
        name = fqn.split(".")[-1]
        prompts.append(f"What does `{name}` do in {module}? Show signature and cite path:line.")
    # Unique and bounded
    uniq: List[str] = []
    for p in prompts:
        if p not in uniq:


# citation regex check, pytest test runner

from typing import Dict, List, Any
import subprocess
from ..code_graph import CodeGraph
import os

def verify_with_tests(g: CodeGraph, module: str, *, repo_root: str, env: Dict[str, str]) -> bool:
    """Run mapped tests for a module if available; return True if all selected tests pass.

    Uses CodeGraph's tests mapping (best-effort). If no tests mapped, returns True as a soft pass.
    """
    nodes = g.pytest_nodes_by_module.get(module, [])
    # Also include module-level mapping
    if not nodes:
        nodes = g.tests_for_module(module)
    if not nodes:
        return True  # no tests to run; accept
    # Build pytest command
    cmd = ["pytest", "-q"]


# Example-specific utilities for the Python repo backend


# question-aware selection, zoom-by-symbol, function-first candidate collection, name-match helpers
import os
import re
from typing import Dict, List, Optional, Tuple

from ..code_graph import CodeGraph

def _re_escape(s: str) -> str:
    try:
        import re as _re
        return _re.escape(s)
    except Exception:
        return s

def modules_from_symbols(repo_root: str, seeds: List[str], *, radius: int = 1, top_k: int = 8, ignore: Optional[List[str]] = None) -> Tuple[List[str], List[str]]:
    g = CodeGraph.load_or_build(repo_root, ignore=[s for s in (ignore or []) if s])
    modules_set: Dict[str, bool] = {}
    files_set: Dict[str, bool] = {}
    for s in seeds:


from __future__ import annotations

import os
import json
import time
import hashlib
import subprocess
import sys
from typing import Dict, List, Tuple, Optional, Set

from ..code_graph import CodeGraph
from .prompts import build_prompts_for_module
from .verify import verify_with_tests, extract_citations
from modules.embedding import build_subgraph_embedding_from_graph
from modules.adapter import (
    save_npz,
    generate_lora_from_embedding,
)
from model.inspect import detect_target_shapes_from_model_full, detect_target_shapes_from_model


import os
import json
import argparse
from typing import Dict, Any, List

import torch

@torch.no_grad()
def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--prompt", default="Hello")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--dtype", default="bfloat16")
    p.add_argument("--max-layers", type=int, default=8, help="Limit layers to trace for brevity")
    p.add_argument("--save-logits", action="store_true", help="Save HF/local logits to disk to minimize RAM/VRAM use")
    p.add_argument("--out-dir", default="", help="Directory to write saved tensors (defaults to cache-dir/trace)")
    args = p.parse_args()


import os
import json
from typing import List, Dict

def main() -> None:
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True, help="HF model id or local snapshot path")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--samples", type=int, default=0, help="Sample count from a default suite (0 uses all)")
    p.add_argument("--backend", default="local", choices=["hf","local","pure"], help="Which local tokenizer backend to compare against HF")
    args = p.parse_args()

    from model.hf_snapshot import ensure_snapshot
    snap = ensure_snapshot(args.model, args.cache_dir)

    # HF reference tokenizer
    try:
        from transformers import AutoTokenizer  # type: ignore
    except Exception as e:


import os
import json
import argparse
from typing import Dict, List, Tuple

import torch

def _ensure_snapshot(model_id: str, cache_dir: str) -> str:
    if os.path.isdir(model_id) and os.path.isfile(os.path.join(model_id, "config.json")):
        return model_id
    org_name = model_id.strip().split("/")[-2:]
    if len(org_name) == 2:
        org, name = org_name
        dir1 = os.path.join(cache_dir, f"models--{org}--{name}", "snapshots")
        cands = []
        if os.path.isdir(dir1):
            cands.extend([os.path.join(dir1, d) for d in os.listdir(dir1)])
        cands = [p for p in cands if os.path.isfile(os.path.join(p, "model.safetensors.index.json"))]
        if cands:
            cands.sort(key=lambda p: os.path.getmtime(p), reverse=True)


import os
import sys
import json
import gc
import torch
from typing import Tuple
from model.hf_snapshot import ensure_snapshot
from model.runtime_utils import local_logits_last

def main() -> None:
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="meta-llama/Llama-3.1-8B-Instruct")
    p.add_argument("--cache-dir", default="/data/transformer_10/checkpoints")
    p.add_argument("--prompt", default="Hello")
    p.add_argument("--max-tokens", type=int, default=64)
    p.add_argument("--gen-new", type=int, default=32, help="Generate this many new tokens for HF/local side-by-side")
    p.add_argument("--mode", default="hf_guided", choices=["independent", "hf_guided", "single_step", "independent_exact", "two_pass_exact"], help="Verification mode")
    p.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"], help="Computation dtype for both models during verify")
    p.add_argument("--save-logits", action="store_true", help="Save HF/local last-step logits to disk to minimize RAM/VRAM use")
    p.add_argument("--parity-exact", action="store_true", help="Use explicit matmul+float32 softmax attention for exact HF parity")