                    {
                        "node_id": rec.node_id,
                        "level": rec.level,
                        "type": rec.type,
                        "parent_id": rec.parent_id,
                        "name": rec.name,
                        "attributes": rec.attributes,
                    }
                )
                + "\n"
            )

def _write_level_sizes(level_counts: Dict[int, int], out_path: Path) -> None:
    # Exclude root (level 0) from level_sizes; TOLBERT heads start at level 1.
    level_sizes = {level: count for level, count in level_counts.items() if level > 0}
    with out_path.open("w", encoding="utf-8") as f:
        json.dump({"level_sizes": level_sizes}, f, indent=2)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(
        description="Build a joint code+paper Tree-of-Life from per-domain spans and nodes.",
    )
    ap.add_argument("--code_spans", type=str, required=True, help="Code spans JSONL file.")
    ap.add_argument("--code_nodes", type=str, required=True, help="Code nodes JSONL file.")
    ap.add_argument("--paper_spans", type=str, required=True, help="Paper spans JSONL file.")
    ap.add_argument("--paper_nodes", type=str, required=True, help="Paper nodes JSONL file.")

    ap.add_argument(
        "--out_code_spans",
        type=str,
        required=True,
        help="Output path for code spans JSONL with joint node_path.",
    )
    ap.add_argument(
        "--out_paper_spans",
        type=str,
        required=True,
        help="Output path for paper spans JSONL with joint node_path.",
    )
    ap.add_argument(
        "--out_nodes",
        type=str,
        required=True,
        help="Output path for joint nodes JSONL.",
    )
    ap.add_argument(
        "--out_level_sizes",
        type=str,
        required=True,
        help="Output path for joint level_sizes JSON helper JSON.",
    )

    ap.add_argument(
        "--code_domain_name",
        type=str,
        default="Code",
        help="Name for the code domain node at level 1.",
    )
    ap.add_argument(
        "--paper_domain_name",
        type=str,
        default="Papers",
        help="Name for the paper domain node at level 1.",
    )

    return ap.parse_args()

def main() -> None:
    args = parse_args()

    code_nodes_path = Path(args.code_nodes)
    paper_nodes_path = Path(args.paper_nodes)
    code_spans_path = Path(args.code_spans)
    paper_spans_path = Path(args.paper_spans)

    if not code_nodes_path.is_file():
        raise FileNotFoundError(f"code_nodes does not exist or is not a file: {code_nodes_path}")
    if not paper_nodes_path.is_file():
        raise FileNotFoundError(
            f"paper_nodes does not exist or is not a file: {paper_nodes_path}"
        )
    if not code_spans_path.is_file():
        raise FileNotFoundError(f"code_spans does not exist or is not a file: {code_spans_path}")
    if not paper_spans_path.is_file():
        raise FileNotFoundError(


main

        required=True,
        help="Output path for joint level_sizes JSON helper JSON.",
    )

    ap.add_argument(
        "--code_domain_name",
        type=str,
        default="Code",
        help="Name for the code domain node at level 1.",
    )
    ap.add_argument(
        "--paper_domain_name",
        type=str,
        default="Papers",
        help="Name for the paper domain node at level 1.",
    )

    return ap.parse_args()

def main() -> None:
    args = parse_args()

    code_nodes_path = Path(args.code_nodes)
    paper_nodes_path = Path(args.paper_nodes)
    code_spans_path = Path(args.code_spans)
    paper_spans_path = Path(args.paper_spans)

    if not code_nodes_path.is_file():
        raise FileNotFoundError(f"code_nodes does not exist or is not a file: {code_nodes_path}")
    if not paper_nodes_path.is_file():
        raise FileNotFoundError(
            f"paper_nodes does not exist or is not a file: {paper_nodes_path}"
        )
    if not code_spans_path.is_file():
        raise FileNotFoundError(f"code_spans does not exist or is not a file: {code_spans_path}")
    if not paper_spans_path.is_file():
        raise FileNotFoundError(
            f"paper_spans does not exist or is not a file: {paper_spans_path}"
        )

    code_nodes = _load_nodes(code_nodes_path)
    paper_nodes = _load_nodes(paper_nodes_path)

    joint_nodes, code_id_map, paper_id_map, level_counts = _build_joint_ontology(
        code_nodes=code_nodes,
        paper_nodes=paper_nodes,
        code_domain_name=args.code_domain_name,
        paper_domain_name=args.paper_domain_name,
    )

    # Rewrite spans into the joint ID space.
    # Domain node IDs are always 1 (code) and 2 (paper) per _build_joint_ontology.
    code_domain_id = 1
    paper_domain_id = 2
    code_spans_joint = _rewrite_spans(
        code_spans_path,
        domain_node_id=code_domain_id,
        id_map=code_id_map,
    )
    paper_spans_joint = _rewrite_spans(
        paper_spans_path,
        domain_node_id=paper_domain_id,
        id_map=paper_id_map,
    )

    out_code_spans_path = Path(args.out_code_spans)
    out_paper_spans_path = Path(args.out_paper_spans)
    out_nodes_path = Path(args.out_nodes)
    out_level_sizes_path = Path(args.out_level_sizes)

    out_code_spans_path.parent.mkdir(parents=True, exist_ok=True)
    out_paper_spans_path.parent.mkdir(parents=True, exist_ok=True)
    out_nodes_path.parent.mkdir(parents=True, exist_ok=True)
    out_level_sizes_path.parent.mkdir(parents=True, exist_ok=True)

    _write_spans(out_code_spans_path, code_spans_joint)
    _write_spans(out_paper_spans_path, paper_spans_joint)
    _write_nodes_jsonl(joint_nodes, out_nodes_path)
    _write_level_sizes(level_counts, out_level_sizes_path)

if __name__ == "__main__":
    main()


        node_id=code_domain_id,
        level=1,
        type="domain",
        parent_id=root_id,
        name=code_domain_name,
        attributes={"source": "code"},
    )
    joint_nodes[paper_domain_id] = NodeRecord(
        node_id=paper_domain_id,
        level=1,
        type="domain",
        parent_id=root_id,
        name=paper_domain_name,
        attributes={"source": "paper"},
    )
    level_counts[1] = 2

    code_id_map: Dict[int, int] = {}
    paper_id_map: Dict[int, int] = {}

    def _remap_domain(
        *,
        src_nodes: Dict[int, NodeRecord],
        domain_node_id: int,
        is_code: bool,
    ) -> Dict[int, int]:
        nonlocal next_id
        id_map: Dict[int, int] = {}

        # Sort by (level, node_id) for stable assignment.
        for old_id, rec in sorted(src_nodes.items(), key=lambda kv: (kv[1].level, kv[0])):
            if rec.level == 0:
                # Skip local root; we replace it with shared root + domain nodes.
                continue

            new_level = rec.level + 1  # shift deeper by one level
            new_id = next_id
            next_id += 1

            # Determine parent in joint space.
            if rec.parent_id is None or rec.parent_id == 0:
                new_parent = domain_node_id
            else:
                if rec.parent_id not in id_map:
                    raise ValueError(
                        f"Parent node_id {rec.parent_id} for node {old_id} has not been remapped yet."
                    )
                new_parent = id_map[rec.parent_id]

            # Merge attributes and tag source.
            attrs = dict(rec.attributes or {})
            attrs.setdefault("source", "code" if is_code else "paper")

            joint_nodes[new_id] = NodeRecord(
                node_id=new_id,
                level=new_level,
                type=rec.type,
                parent_id=new_parent,
                name=rec.name,
                attributes=attrs,
            )

            id_map[old_id] = new_id
            level_counts[new_level] = level_counts.get(new_level, 0) + 1

        return id_map

    code_id_map = _remap_domain(src_nodes=code_nodes, domain_node_id=code_domain_id, is_code=True)
    paper_id_map = _remap_domain(
        src_nodes=paper_nodes, domain_node_id=paper_domain_id, is_code=False
    )

    return joint_nodes, code_id_map, paper_id_map, level_counts

def _rewrite_spans(
    spans_path: Path,
    *,
    domain_node_id: int,
    id_map: Dict[int, int],
) -> List[Dict[str, object]]:
    """
    Rewrite node_path for spans from a single domain into the joint ID space.
    """
    out: List[Dict[str, object]] = []
    root_id = 0


"""
Evaluate a flat (leaf-level) baseline classifier on a labeled spans_file.

This computes standard leaf-level accuracy for the "BERT-flat" style baselines
trained with scripts/train_flat_baseline.py. It does NOT reconstruct full
hierarchical paths; it simply evaluates how well the model predicts the leaf
node id (last element of node_path).

Example:

  python -m scripts.eval_flat_baseline \\
      --config configs/codehierarchy_example.yaml \\
      --checkpoint-dir checkpoints/codehierarchy_bert_flat \\
      --spans-file /data/tolbert/data/codehierarchy/spans_test.jsonl
"""

from __future__ import annotations

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

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    PreTrainedTokenizerBase,
)
import os

from tolbert.config import load_tolbert_config

class FlatLeafEvalDataset(Dataset):
    """
    Identical to FlatLeafDataset but without any training-specific behavior.
    """

    def __init__(
        self,
        spans_file: str,
        tokenizer: PreTrainedTokenizerBase,
        max_length: int = 256,
    ) -> None:
        self.spans_file = spans_file
        self.tokenizer = tokenizer
        self.max_length = max_length

        self._records: List[Dict[str, Any]] = []
        with open(spans_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if "text" not in obj or "node_path" not in obj:
                    continue
                if not obj["node_path"]:
                    continue
                self._records.append(obj)

    def __len__(self) -> int:
        return len(self._records)

    def _tokenize(self, text: str) -> Dict[str, torch.Tensor]:
        enc = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.tensor(label, dtype=torch.long),
        }

def collate_flat_eval_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    input_ids = torch.stack([b["input_ids"] for b in batch], dim=0)
    attention_mask = torch.stack([b["attention_mask"] for b in batch], dim=0)
    labels = torch.stack([b["labels"] for b in batch], dim=0)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate flat (leaf-level) baseline classifier.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Config with max_length / batch_size; level_sizes is used to sanity-check num_labels.",
    )
    ap.add_argument(
        "--checkpoint-dir",
        type=str,
        required=True,
        help="Directory containing the fine-tuned baseline model (save_pretrained output).",
    )
    ap.add_argument(
        "--spans-file",
        type=str,
        required=True,
        help="Labeled spans JSONL file for evaluation (e.g., *_test.jsonl).",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use.",
    )
    return ap.parse_args()

def main() -> None:
    args = parse_args()
    cfg = load_tolbert_config(args.config)

    device = torch.device(args.device)

    spans_path = Path(args.spans_file)
    if not spans_path.exists():
        raise FileNotFoundError(f"spans_file not found: {spans_path}")

    checkpoint_dir = Path(args.checkpoint_dir)
    if not checkpoint_dir.exists():
        raise FileNotFoundError(f"checkpoint_dir not found: {checkpoint_dir}")

    # Leaf class count from config (for sanity checking).
    level_sizes: Dict[int, int] = cfg["level_sizes"]
    leaf_level = max(level_sizes.keys())
    num_labels_expected = level_sizes[leaf_level]

    tokenizer = AutoTokenizer.from_pretrained(
        checkpoint_dir,
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        checkpoint_dir,
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    if model.num_labels != num_labels_expected:
        print(
            f"Warning: model.num_labels={model.num_labels} "
            f"but config expects {num_labels_expected} leaf classes."
        )

    model.to(device)
    model.eval()

    dataset = FlatLeafEvalDataset(
        spans_file=str(spans_path),
        tokenizer=tokenizer,
        max_length=cfg.get("max_length", 256),
    )
    dataloader = DataLoader(
        dataset,
        batch_size=cfg.get("batch_size", 64),
        shuffle=False,
        num_workers=cfg.get("num_workers", 0),
        collate_fn=collate_flat_eval_batch,
    )

    correct = 0
    total = 0

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            logits = outputs.logits
            preds = logits.argmax(dim=-1)

            correct += (preds == labels).sum().item()
            total += labels.size(0)

    acc = correct / max(1, total)
    print("=== Flat Leaf-Level Classification Evaluation ===")
    print(f"Accuracy (leaf node id): {acc:.4f} (n={total})")

if __name__ == "__main__":
    main()


FlatLeafEvalDataset

from __future__ import annotations

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

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    PreTrainedTokenizerBase,
)
import os

from tolbert.config import load_tolbert_config

class FlatLeafEvalDataset(Dataset):
    """
    Identical to FlatLeafDataset but without any training-specific behavior.
    """

    def __init__(
        self,
        spans_file: str,
        tokenizer: PreTrainedTokenizerBase,
        max_length: int = 256,
    ) -> None:
        self.spans_file = spans_file
        self.tokenizer = tokenizer
        self.max_length = max_length

        self._records: List[Dict[str, Any]] = []
        with open(spans_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if "text" not in obj or "node_path" not in obj:
                    continue
                if not obj["node_path"]:
                    continue
                self._records.append(obj)

    def __len__(self) -> int:
        return len(self._records)

    def _tokenize(self, text: str) -> Dict[str, torch.Tensor]:
        enc = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.tensor(label, dtype=torch.long),
        }

def collate_flat_eval_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    input_ids = torch.stack([b["input_ids"] for b in batch], dim=0)
    attention_mask = torch.stack([b["attention_mask"] for b in batch], dim=0)
    labels = torch.stack([b["labels"] for b in batch], dim=0)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate flat (leaf-level) baseline classifier.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Config with max_length / batch_size; level_sizes is used to sanity-check num_labels.",


            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.tensor(label, dtype=torch.long),
        }

def collate_flat_eval_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    input_ids = torch.stack([b["input_ids"] for b in batch], dim=0)
    attention_mask = torch.stack([b["attention_mask"] for b in batch], dim=0)
    labels = torch.stack([b["labels"] for b in batch], dim=0)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate flat (leaf-level) baseline classifier.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Config with max_length / batch_size; level_sizes is used to sanity-check num_labels.",
    )
    ap.add_argument(
        "--checkpoint-dir",
        type=str,
        required=True,
        help="Directory containing the fine-tuned baseline model (save_pretrained output).",
    )
    ap.add_argument(
        "--spans-file",
        type=str,
        required=True,


        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.tensor(label, dtype=torch.long),
        }

def collate_flat_eval_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    input_ids = torch.stack([b["input_ids"] for b in batch], dim=0)
    attention_mask = torch.stack([b["attention_mask"] for b in batch], dim=0)
    labels = torch.stack([b["labels"] for b in batch], dim=0)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate flat (leaf-level) baseline classifier.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Config with max_length / batch_size; level_sizes is used to sanity-check num_labels.",
    )
    ap.add_argument(
        "--checkpoint-dir",
        type=str,
        required=True,
        help="Directory containing the fine-tuned baseline model (save_pretrained output).",
    )
    ap.add_argument(
        "--spans-file",
        type=str,
        required=True,
        help="Labeled spans JSONL file for evaluation (e.g., *_test.jsonl).",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use.",
    )
    return ap.parse_args()

def main() -> None:
    args = parse_args()
    cfg = load_tolbert_config(args.config)

    device = torch.device(args.device)

    spans_path = Path(args.spans_file)
    if not spans_path.exists():
        raise FileNotFoundError(f"spans_file not found: {spans_path}")

    checkpoint_dir = Path(args.checkpoint_dir)
    if not checkpoint_dir.exists():
        raise FileNotFoundError(f"checkpoint_dir not found: {checkpoint_dir}")

    # Leaf class count from config (for sanity checking).
    level_sizes: Dict[int, int] = cfg["level_sizes"]
    leaf_level = max(level_sizes.keys())
    num_labels_expected = level_sizes[leaf_level]


main

        "--checkpoint-dir",
        type=str,
        required=True,
        help="Directory containing the fine-tuned baseline model (save_pretrained output).",
    )
    ap.add_argument(
        "--spans-file",
        type=str,
        required=True,
        help="Labeled spans JSONL file for evaluation (e.g., *_test.jsonl).",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use.",
    )
    return ap.parse_args()

def main() -> None:
    args = parse_args()
    cfg = load_tolbert_config(args.config)

    device = torch.device(args.device)

    spans_path = Path(args.spans_file)
    if not spans_path.exists():
        raise FileNotFoundError(f"spans_file not found: {spans_path}")

    checkpoint_dir = Path(args.checkpoint_dir)
    if not checkpoint_dir.exists():
        raise FileNotFoundError(f"checkpoint_dir not found: {checkpoint_dir}")

    # Leaf class count from config (for sanity checking).
    level_sizes: Dict[int, int] = cfg["level_sizes"]
    leaf_level = max(level_sizes.keys())
    num_labels_expected = level_sizes[leaf_level]

    tokenizer = AutoTokenizer.from_pretrained(
        checkpoint_dir,
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        checkpoint_dir,
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    if model.num_labels != num_labels_expected:
        print(
            f"Warning: model.num_labels={model.num_labels} "
            f"but config expects {num_labels_expected} leaf classes."
        )

    model.to(device)
    model.eval()

    dataset = FlatLeafEvalDataset(
        spans_file=str(spans_path),
        tokenizer=tokenizer,
        max_length=cfg.get("max_length", 256),
    )
    dataloader = DataLoader(
        dataset,
        batch_size=cfg.get("batch_size", 64),
        shuffle=False,
        num_workers=cfg.get("num_workers", 0),
        collate_fn=collate_flat_eval_batch,
    )

    correct = 0
    total = 0

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
            )
            logits = outputs.logits
            preds = logits.argmax(dim=-1)

            correct += (preds == labels).sum().item()
            total += labels.size(0)

    acc = correct / max(1, total)
    print("=== Flat Leaf-Level Classification Evaluation ===")
    print(f"Accuracy (leaf node id): {acc:.4f} (n={total})")

if __name__ == "__main__":
    main()


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

import torch
from torch.utils.data import DataLoader, Dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    PreTrainedTokenizerBase,
)
import os

from tolbert.config import load_tolbert_config

class FlatLeafEvalDataset(Dataset):
    """
    Identical to FlatLeafDataset but without any training-specific behavior.
    """

    def __init__(
        self,
        spans_file: str,
        tokenizer: PreTrainedTokenizerBase,
        max_length: int = 256,
    ) -> None:
        self.spans_file = spans_file
        self.tokenizer = tokenizer
        self.max_length = max_length

        self._records: List[Dict[str, Any]] = []
        with open(spans_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if "text" not in obj or "node_path" not in obj:
                    continue
                if not obj["node_path"]:
                    continue
                self._records.append(obj)

    def __len__(self) -> int:
        return len(self._records)

    def _tokenize(self, text: str) -> Dict[str, torch.Tensor]:
        enc = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]


        tokenizer: PreTrainedTokenizerBase,
        max_length: int = 256,
    ) -> None:
        self.spans_file = spans_file
        self.tokenizer = tokenizer
        self.max_length = max_length

        self._records: List[Dict[str, Any]] = []
        with open(spans_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if "text" not in obj or "node_path" not in obj:
                    continue
                if not obj["node_path"]:
                    continue
                self._records.append(obj)

    def __len__(self) -> int:
        return len(self._records)

    def _tokenize(self, text: str) -> Dict[str, torch.Tensor]:
        enc = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,


        self.spans_file = spans_file
        self.tokenizer = tokenizer
        self.max_length = max_length

        self._records: List[Dict[str, Any]] = []
        with open(spans_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                obj = json.loads(line)
                if "text" not in obj or "node_path" not in obj:
                    continue
                if not obj["node_path"]:
                    continue
                self._records.append(obj)

    def __len__(self) -> int:
        return len(self._records)

    def _tokenize(self, text: str) -> Dict[str, torch.Tensor]:
        enc = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.tensor(label, dtype=torch.long),
        }

def collate_flat_eval_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    input_ids = torch.stack([b["input_ids"] for b in batch], dim=0)
    attention_mask = torch.stack([b["attention_mask"] for b in batch], dim=0)
    labels = torch.stack([b["labels"] for b in batch], dim=0)
    return {


                obj = json.loads(line)
                if "text" not in obj or "node_path" not in obj:
                    continue
                if not obj["node_path"]:
                    continue
                self._records.append(obj)

    def __len__(self) -> int:
        return len(self._records)

    def _tokenize(self, text: str) -> Dict[str, torch.Tensor]:
        enc = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
        )
        return {k: v.squeeze(0) for k, v in enc.items()}

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        rec = self._records[idx]
        tokens = self._tokenize(rec["text"])
        input_ids = tokens["input_ids"]
        attention_mask = tokens.get("attention_mask", torch.ones_like(input_ids))
        node_path: List[int] = rec["node_path"]
        label = int(node_path[-1])
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.tensor(label, dtype=torch.long),
        }

def collate_flat_eval_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    input_ids = torch.stack([b["input_ids"] for b in batch], dim=0)
    attention_mask = torch.stack([b["attention_mask"] for b in batch], dim=0)
    labels = torch.stack([b["labels"] for b in batch], dim=0)
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels,
    }

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate flat (leaf-level) baseline classifier.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Config with max_length / batch_size; level_sizes is used to sanity-check num_labels.",


"""
Minimal training skeleton for TOLBERT **with curriculum and multi-domain support**.

This script wires together:
  - config loading
  - tokenizer + dataset(s)
  - model + optimizer
  - a training loop with optional curriculum over hierarchical / path / contrastive losses

It assumes you have prepared one or more `spans_file`(s) as described in `docs/tree_of_life.md`.
"""

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

import torch
from torch.utils.data import DataLoader, ConcatDataset
from transformers import AutoTokenizer
import os

from tolbert.config import load_tolbert_config
from tolbert.data import TreeOfLifeDataset, collate_tree_of_life_batch
from tolbert.losses import tree_contrastive_loss
from tolbert.modeling import TOLBERT, TOLBERTConfig

def build_model(cfg: Dict[str, Any]) -> TOLBERT:
    """
    Build a TOLBERT model from config.

    Notes:
      - `lambda_hier` and `lambda_path` are used as *default* weights inside
        the model when no curriculum is configured.
      - When a curriculum is enabled, per-step weights are applied in the
        training loop based on `loss_components`, and these config values act
        only as fallbacks.
    """
    model_cfg = TOLBERTConfig(
        base_model_name=cfg["base_model_name"],
        level_sizes=cfg["level_sizes"],
        proj_dim=cfg.get("proj_dim", 256),
        lambda_hier=cfg.get("lambda_hier", 1.0),
        lambda_path=cfg.get("lambda_path", 0.0),
        lambda_contrast=cfg.get("lambda_contrast", 0.0),
    )
    return TOLBERT(model_cfg)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Train a TOLBERT model.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to YAML/JSON config file describing model and training params.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use (default: cuda if available, else cpu).",
    )
    return ap.parse_args()

def _build_dataset(cfg: Dict[str, Any], tokenizer: Any) -> torch.utils.data.Dataset:
    """
    Build a (possibly multi-domain) dataset from config.

    Supported config patterns:
      - `spans_file: /path/to/spans.jsonl`
      - `spans_files: [ /path/to/code.jsonl, /path/to/papers.jsonl, ... ]`

    In the multi-file case we simply concatenate datasets; mixed-domain
    batches then come from the union of underlying corpora, as described
    in the paper.
    """
    max_length = cfg.get("max_length", 256)
    mask_probability = cfg.get("mask_probability", 0.15)

    if "spans_files" in cfg and cfg["spans_files"] is not None:
        spans_files_cfg = cfg["spans_files"]
        if isinstance(spans_files_cfg, str):
            spans_files = [spans_files_cfg]
        else:
            spans_files = list(spans_files_cfg)

        datasets = []
        for spans_path in spans_files:
            if not Path(spans_path).exists():
                raise FileNotFoundError(f"spans_file not found: {spans_path}")
            datasets.append(
                TreeOfLifeDataset(
                    spans_file=spans_path,
                    tokenizer=tokenizer,
                    max_length=max_length,
                    mask_probability=mask_probability,
                )
            )

        if len(datasets) == 1:
            return datasets[0]
        return ConcatDataset(datasets)

    # Fallback: single-file setup (backwards-compatible).
    spans_file = cfg["spans_file"]
    if not Path(spans_file).exists():
        raise FileNotFoundError(f"spans_file not found: {spans_file}")

    return TreeOfLifeDataset(
        spans_file=spans_file,
        tokenizer=tokenizer,
        max_length=max_length,
        mask_probability=mask_probability,
    )

def _get_curriculum_stage(
    curriculum_cfg: Optional[Dict[str, Any]],
    global_step: int,
) -> Optional[Dict[str, Any]]:
    """
    Retrieve the active curriculum stage for the given step.

    Expected config shape:
      curriculum:
        enabled: true
        stages:
          - name: warmup
            start_step: 0
            end_step: 10000
            max_supervised_level: 2
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.0
          - name: deep
            start_step: 10000
            end_step: 20000
            max_supervised_level: 3
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.0
          - name: contrastive
            start_step: 20000
            end_step: 30000
            max_supervised_level: 3
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.05
          - name: path
            start_step: 30000
            end_step: null    # or omit to mean "until end"
            max_supervised_level: 3
            lambda_hier: 1.0
            lambda_path: 0.05
            lambda_contrast: 0.05

    If `curriculum.enabled` is false or missing, returns None.
    """
    if not curriculum_cfg or not curriculum_cfg.get("enabled", False):
        return None

    stages = curriculum_cfg.get("stages") or []
    if not stages:
        return None

    for stage in stages:
        start = int(stage.get("start_step", 0))
        end_raw = stage.get("end_step", None)
        end: Optional[int] = None if end_raw is None else int(end_raw)

        if global_step < start:
            continue
        if end is not None and global_step >= end:
            continue
        return stage

    # If we fall past all defined ranges, use the last stage as a default.
    return stages[-1]

def main() -> None:
    args = parse_args()
    cfg = load_tolbert_config(args.config)

    device = torch.device(args.device)

    # Tokenizer and dataset(s)
    tokenizer = AutoTokenizer.from_pretrained(
        cfg["base_model_name"],
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    dataset = _build_dataset(cfg, tokenizer)

    dataloader = DataLoader(
        dataset,
        batch_size=cfg.get("batch_size", 64),
        shuffle=True,
        num_workers=cfg.get("num_workers", 0),
        collate_fn=collate_tree_of_life_batch,
    )

    # Model and optimizer
    model = build_model(cfg).to(device)
    model.train()

    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.get("lr", 1e-4))

    num_epochs = cfg.get("num_epochs", 1)
    contrast_temp = cfg.get("contrast_temperature", 0.07)
    curriculum_cfg: Optional[Dict[str, Any]] = cfg.get("curriculum")

    # Fallback static weights when curriculum is disabled.
    static_lambda_hier = float(cfg.get("lambda_hier", 1.0))
    static_lambda_path = float(cfg.get("lambda_path", 0.0))
    static_lambda_contrast = float(cfg.get("lambda_contrast", 0.0))
    use_contrastive_static = static_lambda_contrast > 0.0

    global_step = 0

    for epoch in range(num_epochs):
        for step, batch in enumerate(dataloader, start=1):
            global_step += 1
            optimizer.zero_grad()

            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels_mlm = batch["labels_mlm"].to(device)

            # Move level targets to device and optionally trim by curriculum.
            level_targets = {
                level: targets.to(device) for level, targets in batch["level_targets"].items()
            }

            stage = _get_curriculum_stage(curriculum_cfg, global_step)

            # Determine per-step hyperparameters (MLM is always on).
            if stage is not None:
                max_level = stage.get("max_supervised_level")
                if max_level is not None:
                    max_level_int = int(max_level)
                    level_targets = {
                        level: tgt for level, tgt in level_targets.items() if level <= max_level_int
                    }

                lambda_hier = float(stage.get("lambda_hier", static_lambda_hier))
                lambda_path = float(stage.get("lambda_path", 0.0))
                lambda_contrast = float(stage.get("lambda_contrast", 0.0))
            else:
                lambda_hier = static_lambda_hier
                lambda_path = static_lambda_path
                lambda_contrast = static_lambda_contrast

            # Only pass paths into the model if we actually intend to use path loss.
            paths_for_model = batch.get("paths") if (lambda_path > 0.0 and "paths" in batch) else None

            out = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels_mlm=labels_mlm,
                level_targets=level_targets,
                paths=paths_for_model,
            )

            loss_components = out["loss_components"]
            mlm_loss = loss_components.get("mlm")
            hier_loss = loss_components.get("hier")
            path_loss = loss_components.get("path")

            # Aggregate loss according to either curriculum or static weights.
            if mlm_loss is None:
                raise RuntimeError("MLM loss is expected but missing from model outputs.")

            loss = mlm_loss

            if hier_loss is not None and lambda_hier != 0.0:
                loss = loss + lambda_hier * hier_loss

            if path_loss is not None and lambda_path != 0.0:
                loss = loss + lambda_path * path_loss

            # Contrastive loss: tree-aware supervised contrastive over paths.
            use_contrastive = lambda_contrast > 0.0
            if use_contrastive and "paths" in batch:
                contrast_loss = tree_contrastive_loss(
                    embeddings=out["proj"],
                    paths=batch["paths"],
                    temperature=contrast_temp,
                )
                loss = loss + lambda_contrast * contrast_loss

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=cfg.get("grad_clip", 1.0))
            optimizer.step()

            if step % cfg.get("log_every", 50) == 0:
                loss_items = {k: v.item() for k, v in loss_components.items()}
                print(
                    f"[epoch {epoch+1} step {step} global_step {global_step}] "
                    f"loss={loss.item():.4f} components={loss_items}"
                )

        # Simple checkpointing at end of each epoch
        out_dir = Path(cfg.get("output_dir", "checkpoints"))
        out_dir.mkdir(parents=True, exist_ok=True)
        ckpt_path = out_dir / f"tolbert_epoch{epoch+1}.pt"
        torch.save(model.state_dict(), ckpt_path)
        print(f"Saved checkpoint to {ckpt_path}")

if __name__ == "__main__":
    main()


  - a training loop with optional curriculum over hierarchical / path / contrastive losses

It assumes you have prepared one or more `spans_file`(s) as described in `docs/tree_of_life.md`.
"""

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

import torch
from torch.utils.data import DataLoader, ConcatDataset
from transformers import AutoTokenizer
import os

from tolbert.config import load_tolbert_config
from tolbert.data import TreeOfLifeDataset, collate_tree_of_life_batch
from tolbert.losses import tree_contrastive_loss
from tolbert.modeling import TOLBERT, TOLBERTConfig

def build_model(cfg: Dict[str, Any]) -> TOLBERT:
    """
    Build a TOLBERT model from config.

    Notes:
      - `lambda_hier` and `lambda_path` are used as *default* weights inside
        the model when no curriculum is configured.
      - When a curriculum is enabled, per-step weights are applied in the
        training loop based on `loss_components`, and these config values act
        only as fallbacks.
    """
    model_cfg = TOLBERTConfig(
        base_model_name=cfg["base_model_name"],
        level_sizes=cfg["level_sizes"],
        proj_dim=cfg.get("proj_dim", 256),
        lambda_hier=cfg.get("lambda_hier", 1.0),
        lambda_path=cfg.get("lambda_path", 0.0),
        lambda_contrast=cfg.get("lambda_contrast", 0.0),
    )
    return TOLBERT(model_cfg)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Train a TOLBERT model.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to YAML/JSON config file describing model and training params.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use (default: cuda if available, else cpu).",
    )
    return ap.parse_args()

def _build_dataset(cfg: Dict[str, Any], tokenizer: Any) -> torch.utils.data.Dataset:


    Build a TOLBERT model from config.

    Notes:
      - `lambda_hier` and `lambda_path` are used as *default* weights inside
        the model when no curriculum is configured.
      - When a curriculum is enabled, per-step weights are applied in the
        training loop based on `loss_components`, and these config values act
        only as fallbacks.
    """
    model_cfg = TOLBERTConfig(
        base_model_name=cfg["base_model_name"],
        level_sizes=cfg["level_sizes"],
        proj_dim=cfg.get("proj_dim", 256),
        lambda_hier=cfg.get("lambda_hier", 1.0),
        lambda_path=cfg.get("lambda_path", 0.0),
        lambda_contrast=cfg.get("lambda_contrast", 0.0),
    )
    return TOLBERT(model_cfg)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Train a TOLBERT model.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to YAML/JSON config file describing model and training params.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use (default: cuda if available, else cpu).",
    )
    return ap.parse_args()

def _build_dataset(cfg: Dict[str, Any], tokenizer: Any) -> torch.utils.data.Dataset:
    """
    Build a (possibly multi-domain) dataset from config.

    Supported config patterns:
      - `spans_file: /path/to/spans.jsonl`
      - `spans_files: [ /path/to/code.jsonl, /path/to/papers.jsonl, ... ]`

    In the multi-file case we simply concatenate datasets; mixed-domain
    batches then come from the union of underlying corpora, as described
    in the paper.
    """
    max_length = cfg.get("max_length", 256)
    mask_probability = cfg.get("mask_probability", 0.15)

    if "spans_files" in cfg and cfg["spans_files"] is not None:
        spans_files_cfg = cfg["spans_files"]
        if isinstance(spans_files_cfg, str):


    return TOLBERT(model_cfg)

def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Train a TOLBERT model.")
    ap.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to YAML/JSON config file describing model and training params.",
    )
    ap.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device to use (default: cuda if available, else cpu).",
    )
    return ap.parse_args()

def _build_dataset(cfg: Dict[str, Any], tokenizer: Any) -> torch.utils.data.Dataset:
    """
    Build a (possibly multi-domain) dataset from config.

    Supported config patterns:
      - `spans_file: /path/to/spans.jsonl`
      - `spans_files: [ /path/to/code.jsonl, /path/to/papers.jsonl, ... ]`

    In the multi-file case we simply concatenate datasets; mixed-domain
    batches then come from the union of underlying corpora, as described
    in the paper.
    """
    max_length = cfg.get("max_length", 256)
    mask_probability = cfg.get("mask_probability", 0.15)

    if "spans_files" in cfg and cfg["spans_files"] is not None:
        spans_files_cfg = cfg["spans_files"]
        if isinstance(spans_files_cfg, str):
            spans_files = [spans_files_cfg]
        else:
            spans_files = list(spans_files_cfg)

        datasets = []
        for spans_path in spans_files:
            if not Path(spans_path).exists():
                raise FileNotFoundError(f"spans_file not found: {spans_path}")
            datasets.append(
                TreeOfLifeDataset(
                    spans_file=spans_path,
                    tokenizer=tokenizer,
                    max_length=max_length,
                    mask_probability=mask_probability,
                )
            )

        if len(datasets) == 1:
            return datasets[0]
        return ConcatDataset(datasets)

    # Fallback: single-file setup (backwards-compatible).
    spans_file = cfg["spans_file"]
    if not Path(spans_file).exists():
        raise FileNotFoundError(f"spans_file not found: {spans_file}")

    return TreeOfLifeDataset(
        spans_file=spans_file,
        tokenizer=tokenizer,
        max_length=max_length,
        mask_probability=mask_probability,
    )

def _get_curriculum_stage(
    curriculum_cfg: Optional[Dict[str, Any]],
    global_step: int,
) -> Optional[Dict[str, Any]]:
    """
    Retrieve the active curriculum stage for the given step.

    Expected config shape:
      curriculum:
        enabled: true
        stages:
          - name: warmup
            start_step: 0
            end_step: 10000
            max_supervised_level: 2
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.0


                )
            )

        if len(datasets) == 1:
            return datasets[0]
        return ConcatDataset(datasets)

    # Fallback: single-file setup (backwards-compatible).
    spans_file = cfg["spans_file"]
    if not Path(spans_file).exists():
        raise FileNotFoundError(f"spans_file not found: {spans_file}")

    return TreeOfLifeDataset(
        spans_file=spans_file,
        tokenizer=tokenizer,
        max_length=max_length,
        mask_probability=mask_probability,
    )

def _get_curriculum_stage(
    curriculum_cfg: Optional[Dict[str, Any]],
    global_step: int,
) -> Optional[Dict[str, Any]]:
    """
    Retrieve the active curriculum stage for the given step.

    Expected config shape:
      curriculum:
        enabled: true
        stages:
          - name: warmup
            start_step: 0
            end_step: 10000
            max_supervised_level: 2
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.0
          - name: deep
            start_step: 10000
            end_step: 20000
            max_supervised_level: 3
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.0
          - name: contrastive
            start_step: 20000
            end_step: 30000
            max_supervised_level: 3
            lambda_hier: 1.0
            lambda_path: 0.0
            lambda_contrast: 0.05
          - name: path
            start_step: 30000
            end_step: null    # or omit to mean "until end"
            max_supervised_level: 3
            lambda_hier: 1.0
            lambda_path: 0.05
            lambda_contrast: 0.05

    If `curriculum.enabled` is false or missing, returns None.
    """
    if not curriculum_cfg or not curriculum_cfg.get("enabled", False):
        return None

    stages = curriculum_cfg.get("stages") or []
    if not stages:
        return None

    for stage in stages:
        start = int(stage.get("start_step", 0))
        end_raw = stage.get("end_step", None)
        end: Optional[int] = None if end_raw is None else int(end_raw)

        if global_step < start:
            continue
        if end is not None and global_step >= end:
            continue
        return stage

    # If we fall past all defined ranges, use the last stage as a default.
    return stages[-1]

def main() -> None:
    args = parse_args()
    cfg = load_tolbert_config(args.config)

    device = torch.device(args.device)

    # Tokenizer and dataset(s)
    tokenizer = AutoTokenizer.from_pretrained(
        cfg["base_model_name"],
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    dataset = _build_dataset(cfg, tokenizer)

    dataloader = DataLoader(
        dataset,
        batch_size=cfg.get("batch_size", 64),
        shuffle=True,
        num_workers=cfg.get("num_workers", 0),


    stages = curriculum_cfg.get("stages") or []
    if not stages:
        return None

    for stage in stages:
        start = int(stage.get("start_step", 0))
        end_raw = stage.get("end_step", None)
        end: Optional[int] = None if end_raw is None else int(end_raw)

        if global_step < start:
            continue
        if end is not None and global_step >= end:
            continue
        return stage

    # If we fall past all defined ranges, use the last stage as a default.
    return stages[-1]

def main() -> None:
    args = parse_args()
    cfg = load_tolbert_config(args.config)

    device = torch.device(args.device)

    # Tokenizer and dataset(s)
    tokenizer = AutoTokenizer.from_pretrained(
        cfg["base_model_name"],
        cache_dir="/data/checkpoints/",  # noqa: E501
    )
    dataset = _build_dataset(cfg, tokenizer)

    dataloader = DataLoader(
        dataset,
        batch_size=cfg.get("batch_size", 64),
        shuffle=True,
        num_workers=cfg.get("num_workers", 0),
        collate_fn=collate_tree_of_life_batch,
    )

    # Model and optimizer
    model = build_model(cfg).to(device)
    model.train()

    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.get("lr", 1e-4))

    num_epochs = cfg.get("num_epochs", 1)
    contrast_temp = cfg.get("contrast_temperature", 0.07)
    curriculum_cfg: Optional[Dict[str, Any]] = cfg.get("curriculum")

    # Fallback static weights when curriculum is disabled.
    static_lambda_hier = float(cfg.get("lambda_hier", 1.0))
    static_lambda_path = float(cfg.get("lambda_path", 0.0))
    static_lambda_contrast = float(cfg.get("lambda_contrast", 0.0))
    use_contrastive_static = static_lambda_contrast > 0.0

    global_step = 0

    for epoch in range(num_epochs):
        for step, batch in enumerate(dataloader, start=1):
            global_step += 1
            optimizer.zero_grad()

            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels_mlm = batch["labels_mlm"].to(device)

            # Move level targets to device and optionally trim by curriculum.
            level_targets = {
                level: targets.to(device) for level, targets in batch["level_targets"].items()
            }

            stage = _get_curriculum_stage(curriculum_cfg, global_step)

            # Determine per-step hyperparameters (MLM is always on).
            if stage is not None:
                max_level = stage.get("max_supervised_level")
                if max_level is not None:
                    max_level_int = int(max_level)
                    level_targets = {
                        level: tgt for level, tgt in level_targets.items() if level <= max_level_int
                    }

                lambda_hier = float(stage.get("lambda_hier", static_lambda_hier))
                lambda_path = float(stage.get("lambda_path", 0.0))
                lambda_contrast = float(stage.get("lambda_contrast", 0.0))
            else:
                lambda_hier = static_lambda_hier
                lambda_path = static_lambda_path
                lambda_contrast = static_lambda_contrast

            # Only pass paths into the model if we actually intend to use path loss.
            paths_for_model = batch.get("paths") if (lambda_path > 0.0 and "paths" in batch) else None

            out = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels_mlm=labels_mlm,
                level_targets=level_targets,
                paths=paths_for_model,
            )

            loss_components = out["loss_components"]
            mlm_loss = loss_components.get("mlm")
            hier_loss = loss_components.get("hier")
            path_loss = loss_components.get("path")

            # Aggregate loss according to either curriculum or static weights.
            if mlm_loss is None:
                raise RuntimeError("MLM loss is expected but missing from model outputs.")

            loss = mlm_loss

            if hier_loss is not None and lambda_hier != 0.0:
                loss = loss + lambda_hier * hier_loss

            if path_loss is not None and lambda_path != 0.0:
                loss = loss + lambda_path * path_loss

            # Contrastive loss: tree-aware supervised contrastive over paths.
            use_contrastive = lambda_contrast > 0.0
            if use_contrastive and "paths" in batch:
                contrast_loss = tree_contrastive_loss(
                    embeddings=out["proj"],
                    paths=batch["paths"],
                    temperature=contrast_temp,
                )
                loss = loss + lambda_contrast * contrast_loss

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=cfg.get("grad_clip", 1.0))
            optimizer.step()

            if step % cfg.get("log_every", 50) == 0:
                loss_items = {k: v.item() for k, v in loss_components.items()}
                print(
                    f"[epoch {epoch+1} step {step} global_step {global_step}] "
                    f"loss={loss.item():.4f} components={loss_items}"
                )

        # Simple checkpointing at end of each epoch
        out_dir = Path(cfg.get("output_dir", "checkpoints"))
        out_dir.mkdir(parents=True, exist_ok=True)
        ckpt_path = out_dir / f"tolbert_epoch{epoch+1}.pt"
        torch.save(model.state_dict(), ckpt_path)
        print(f"Saved checkpoint to {ckpt_path}")

if __name__ == "__main__":
    main()


from __future__ import annotations

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

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

def artifact_uri(program_id: str, rel_path: str) -> str:
    rel = rel_path.replace("\\", "/").lstrip("/")
    return f"program://{program_id}/artifact/{rel}"

def parse_program_uri(uri: str) -> Tuple[str, str, str, Optional[Tuple[int, int]]]:
    m = re.match(r"^program://([^/]+)/([^/]+)/(.+?)(?:#L(\d+)-L(\d+))?$", uri)
    if not m:
        raise ValueError(f"invalid program uri: {uri}")
    pid, kind, res, a, b = m.group(1), m.group(2), m.group(3), m.group(4), m.group(5)
    span = (int(a), int(b)) if (a and b) else None
    return pid, kind, res, span

class RepoGraph(ProgramGraph):
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.repo_root = os.path.abspath(repo_root)
        self.program_id = program_id_for_repo(self.repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self._file_hash: Dict[str, str] = {}
        # Cached graph views
        self._entities: Dict[EntityId, Entity] = {}
        self._edges: List[Edge] = []
        self._built: bool = False
        # Convenience indices
        #   - absolute file path → file-entity id
        #   - symbol name (lowercased) → list of entity ids
        self._file_entity_for_abs: Dict[str, EntityId] = {}
        self._symbols_by_name: Dict[str, List[EntityId]] = {}

    # ProgramGraph: core views
    def entities(self) -> Iterable[Entity]:
        self._ensure_built()
        return self._entities.values()

    def edges(self) -> Iterable[Edge]:
        self._ensure_built()
        return list(self._edges)

    def search_refs(self, token: str) -> Iterable[Tuple[EntityId, Span]]:
        return []

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        if not seeds or radius <= 0:
            return self
        # Generic BFS over current edges view
        adj: Dict[str, List[str]] = {}
        for e in self.edges():
            adj.setdefault(e.src, []).append(e.dst)
            adj.setdefault(e.dst, []).append(e.src)
        cur = set(seeds)
        seen = set(cur)
        for _ in range(max(1, radius)):
            nxt: Set[str] = set()
            for s in list(cur):
                for nb in adj.get(s, []):
                    if nb not in seen:
                        seen.add(nb)
                        nxt.add(nb)
            cur = nxt
        # For now, RepoGraph exposes a single global view; callers that need
        # an actual induced subgraph can post-filter entities/edges.
        return self  # pragma: no cover - view semantics only

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        if kind not in ("artifact", "source"):
            return []
        self._ensure_built()
        out: List[Artifact] = []
        for fp in self._discover_files(self.repo_root, self.ignore_rules):
            rel = os.path.relpath(fp, self.repo_root).replace("\\", "/")
            out.append(
                Artifact(
                    uri=artifact_uri(self.program_id, rel),
                    type="source",
                    hash=self._hash_for(fp),
                    span=None,
                )
            )
        return out

    def resolve(self, uri: str) -> ResolvedAnchor:
        pid, kind, res, span = parse_program_uri(uri)
        if pid != self.program_id:
            raise ValueError(f"program id mismatch: {pid} != {self.program_id}")
        if kind == "artifact":
            abs_fp = os.path.abspath(os.path.join(self.repo_root, res))
            if not os.path.isfile(abs_fp):
                raise FileNotFoundError(f"artifact not found: {abs_fp}")
            a = int(span[0]) if span else 1
            b = int(span[1]) if span else self._safe_count_lines(abs_fp)
            rel = os.path.relpath(abs_fp, self.repo_root).replace("\\", "/")
            return ResolvedAnchor(
                artifact_uri=artifact_uri(self.program_id, rel),
                span=Span(start_line=a, end_line=b),
                hash=self._hash_for(abs_fp),
            )
        # Let subclass handle entity URIs
        return self._resolve_entity_uri(kind, res, span)

    # Hooks for subclasses
    def _resolve_entity_uri(self, kind: str, resource: str, span: Optional[Tuple[int, int]]) -> ResolvedAnchor:
        raise KeyError(f"unrecognized entity uri for kind={kind}, resource={resource}")

    # Build / utilities
    def _ensure_built(self) -> None:
        if self._built:
            return
        self._build_graph()
        self._built = True

    def _build_graph(self) -> None:
        """
        Populate entity and edge sets for this repository.

        This is intentionally language-agnostic at the core:
          - Always creates a repo-level entity.
          - Always creates file-level entities for discovered source artifacts.

        Language-specific structure (imports, calls, packages / modules) is
        layered on top via simple heuristics and optional backends.
        """
        # Repo entity
        repo_eid: EntityId = f"repo:{self.program_id}"
        if repo_eid not in self._entities:
            self._entities[repo_eid] = Entity(
                id=repo_eid,
                kind="repo",
                uri=f"program://{self.program_id}/repo",
                artifact_uri=None,
                span=None,
                labels=["kind:repo"],
                attributes={"root": self.repo_root},
            )

        # Discover files and create file-level entities
        files = self._discover_files(self.repo_root, self.ignore_rules)
        for fp in files:
            rel = os.path.relpath(fp, self.repo_root).replace("\\", "/")
            uri = artifact_uri(self.program_id, rel)
            labels = ["kind:file"] + self._language_labels_for(rel)
            eid: EntityId = uri  # stable id tied to artifact URI
            if eid not in self._entities:
                self._entities[eid] = Entity(
                    id=eid,
                    kind="file",
                    uri=uri,
                    artifact_uri=uri,
                    span=None,
                    labels=labels,
                    attributes={"rel_path": rel},
                )
            self._file_entity_for_abs[os.path.abspath(fp)] = eid
            # Repo "owns" file
            self._edges.append(
                Edge(
                    src=repo_eid,
                    dst=eid,
                    type="owns",
                    attributes={},
                )
            )

        # Language-specific edges and entities.
        self._build_language_edges_and_entities(files)

    def _build_language_edges_and_entities(self, files: List[str]) -> None:
        """
        Add language-specific structure for supported stacks.

        - C / C++: file-level imports from #include relationships.
        - Go: package entities and import edges.
        - Java: imported type entities and import edges.
        - JS / TS: module/file imports from ES modules / require().
        - Python: optional integration via scripts.codegraph_core.CodeGraph
          for module / function / class entities and imports / calls.
        """
        # Best-effort C / C++ / Go / Java / JS import graphs
        for abs_fp in files:
            rel = os.path.relpath(abs_fp, self.repo_root).replace("\\", "/")
            _, ext = os.path.splitext(rel.lower())
            if ext in (".c", ".h", ".cc", ".cpp", ".cxx", ".hpp"):
                self._add_c_includes(abs_fp)
            elif ext == ".go":
                self._add_go_imports(abs_fp)
            elif ext == ".java":
                self._add_java_imports(abs_fp)
            elif ext in (".js", ".jsx", ".mjs", ".ts", ".tsx"):
                self._add_js_like_imports(abs_fp)

        # Optional: richer Python graph via scripts.codegraph_core
        try:
            from scripts.codegraph_core import CodeGraph as _PyCodeGraph  # type: ignore
        except Exception:
            _PyCodeGraph = None  # type: ignore

        if _PyCodeGraph is not None:
            try:
                cg = _PyCodeGraph(self.repo_root, ignore=self.ignore_rules).build()
            except Exception:
                cg = None
            if cg is not None:
                self._ingest_python_codegraph(cg)

        # Non-Python symbol entities + heuristic call graph.
        self._add_non_python_symbols_and_calls(files)

    def _discover_files(self, root: str, ignore: List[str]) -> List[str]:
        out: List[str] = []
        for dirpath, dirnames, filenames in os.walk(root):
            if any(ig and ig in dirpath for ig in ignore):
                continue
            for fn in filenames:
                ap = os.path.abspath(os.path.join(dirpath, fn))
                out.append(ap)
        return out

    def _language_labels_for(self, rel_path: str) -> List[str]:
        """
        Best-effort language tags for a repo-relative path based on extension.

        These are intentionally coarse and are only used to annotate entities
        (e.g., file-level entities) so that downstream tools can filter or
        group by language family when needed.
        """
        _, ext = os.path.splitext(rel_path.lower())
        labels: List[str] = []
        if ext == ".py":
            labels.append("lang:python")
        elif ext in (".c", ".h"):
            labels.append("lang:c")
        elif ext in (".cc", ".cpp", ".cxx", ".hpp"):
            labels.append("lang:cpp")
        elif ext in (".js", ".jsx", ".mjs"):
            labels.append("lang:js")
        elif ext in (".ts", ".tsx"):
            labels.append("lang:ts")
        elif ext == ".go":
            labels.append("lang:go")
        elif ext == ".java":
            labels.append("lang:java")
        elif ext == ".md":
            labels.append("lang:markdown")
        return labels

    def _safe_count_lines(self, abs_file: str) -> int:
        try:
            with open(abs_file, "r", encoding="utf-8", errors="ignore") as fh:
                return sum(1 for _ in fh)
        except Exception:
            return 1

    def _hash_for(self, abs_file: str) -> str:
        if abs_file in self._file_hash:
            return self._file_hash[abs_file]
        try:
            with open(abs_file, "rb") as fh:
                raw = fh.read()
            h = hashlib.sha256(raw).hexdigest()
        except Exception:
            h = ""
        self._file_hash[abs_file] = h
        return h

    # --- Language-specific helpers ---

    def _add_symbol_entity(
        self,
        abs_fp: str,
        name: str,
        kind: str,
        lang_label: str,
        start_line: int,
        end_line: int,
    ) -> Optional[EntityId]:
        abs_fp = os.path.abspath(abs_fp)
        file_eid = self._file_entity_for_abs.get(abs_fp)
        if not file_eid:
            return None
        # Stable id: lang-specific prefix + file rel path + symbol name.
        rel = os.path.relpath(abs_fp, self.repo_root).replace("\\", "/")
        sym_id: EntityId = f"{lang_label}:{rel}:{name}"
        if sym_id in self._entities:
            return sym_id
        file_ent = self._entities[file_eid]
        span = Span(start_line=int(start_line), end_line=int(end_line))
        uri = f"program://{self.program_id}/sym/{lang_label}/{rel}#{name}"
        labels = [f"kind:{kind}", f"lang:{lang_label}"]
        self._entities[sym_id] = Entity(
            id=sym_id,
            kind=kind,
            uri=uri,
            artifact_uri=file_ent.artifact_uri,
            span=span,
            labels=labels,
            attributes={"name": name, "file": rel},
        )
        # file "owns" symbol
        self._edges.append(
            Edge(
                src=file_eid,
                dst=sym_id,
                type="owns",
                attributes={"lang": lang_label},
            )
        )
        self._symbols_by_name.setdefault(name.lower(), []).append(sym_id)
        return sym_id

    def _add_c_includes(self, abs_fp: str) -> None:
        src_eid = self._file_entity_for_abs.get(os.path.abspath(abs_fp))
        if not src_eid:
            return
        try:
            with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                text = fh.read()
        except Exception:
            return
        # Match #include "local.h" and #include <local.h>
        rx = re.compile(r'^\s*#\s*include\s*[<"]([^">]+)[">]', re.MULTILINE)
        for m in rx.finditer(text):
            target = m.group(1).strip()
            if not target:
                continue
            # Resolve relative to current directory; ignore obvious system headers.
            if "/" not in target and "\\" not in target and "." not in target:
                continue
            cand = os.path.abspath(os.path.join(os.path.dirname(abs_fp), target))
            dst_eid = self._file_entity_for_abs.get(cand)
            if not dst_eid:
                continue
            self._edges.append(
                Edge(
                    src=src_eid,
                    dst=dst_eid,
                    type="imports",
                    attributes={"lang": "c_cpp"},
                )
            )

    def _add_go_imports(self, abs_fp: str) -> None:
        src_eid = self._file_entity_for_abs.get(os.path.abspath(abs_fp))
        if not src_eid:
            return
        try:
            with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                lines = fh.readlines()
        except Exception:
            return
        in_block = False
        for ln in lines:
            stripped = ln.strip()
            if stripped.startswith("import "):
                # Single-line: import "pkg/path"
                in_block = "(" in stripped and not stripped.rstrip().endswith(")")
                m = re.search(r'["`](.+?)["`]', stripped)
                if m:
                    pkg = m.group(1).strip()
                    self._add_go_import_edge(src_eid, pkg)
            elif in_block:
                if stripped.startswith(")"):
                    in_block = False
                    continue
                m = re.search(r'["`](.+?)["`]', stripped)
                if m:
                    pkg = m.group(1).strip()
                    self._add_go_import_edge(src_eid, pkg)

    def _add_go_import_edge(self, src_eid: EntityId, pkg: str) -> None:
        if not pkg:
            return
        eid: EntityId = f"go:pkg:{pkg}"
        if eid not in self._entities:
            self._entities[eid] = Entity(
                id=eid,
                kind="package",
                uri=f"program://{self.program_id}/go_pkg/{pkg}",
                artifact_uri=None,
                span=None,
                labels=["kind:package", "lang:go"],
                attributes={"package": pkg},
            )
        self._edges.append(
            Edge(
                src=src_eid,
                dst=eid,
                type="imports",
                    attributes={"lang": "go"},
            )
        )

    def _add_java_imports(self, abs_fp: str) -> None:
        src_eid = self._file_entity_for_abs.get(os.path.abspath(abs_fp))
        if not src_eid:
            return
        try:
            with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                text = fh.read()
        except Exception:
            return
        rx = re.compile(r'^\s*import\s+(static\s+)?([a-zA-Z0-9_.]+)\s*;', re.MULTILINE)
        for m in rx.finditer(text):
            fqn = m.group(2)
            if not fqn:
                continue
            eid: EntityId = f"java:import:{fqn}"
            if eid not in self._entities:
                self._entities[eid] = Entity(
                    id=eid,
                    kind="type",
                    uri=f"program://{self.program_id}/java_type/{fqn}",
                    artifact_uri=None,
                    span=None,
                    labels=["kind:type", "lang:java"],
                    attributes={"fqn": fqn},
                )
            self._edges.append(
                Edge(
                    src=src_eid,
                    dst=eid,
                    type="imports",
                    attributes={"lang": "java"},
                )
            )

    def _add_js_like_imports(self, abs_fp: str) -> None:
        src_eid = self._file_entity_for_abs.get(os.path.abspath(abs_fp))
        if not src_eid:
            return
        try:
            with open(abs_fp, "r", encoding="utf-8", errors="ignore") as fh:
                text = fh.read()
        except Exception:
            return
        rel = os.path.relpath(abs_fp, self.repo_root).replace("\\", "/")
        base_dir = os.path.dirname(os.path.abspath(abs_fp))

        specs: Set[str] = set()
        # import x from 'spec';  import 'spec';
        for m in re.finditer(r'^\s*import\s+(?:.+?\s+from\s+)?[\'"](.+?)[\'"]', text, re.MULTILINE):
            specs.add(m.group(1).strip())
        # require('spec')
        for m in re.finditer(r'require\(\s*[\'"](.+?)[\'"]\s*\)', text):
            specs.add(m.group(1).strip())

        for spec in specs:
            if not spec:
                continue
            if spec.startswith("."):
                # Local module → resolve to file entity if possible.
                dst_eid = self._resolve_local_js_like_spec(base_dir, spec)
                if dst_eid:
                    self._edges.append(
                        Edge(
                            src=src_eid,
                            dst=dst_eid,
                            type="imports",
                            attributes={"lang": "js_ts"},
                        )
                    )
            else:
                # External package
                eid: EntityId = f"js:pkg:{spec}"
                if eid not in self._entities:
                    self._entities[eid] = Entity(
                        id=eid,
                        kind="package",
                        uri=f"program://{self.program_id}/js_pkg/{spec}",
                        artifact_uri=None,
                        span=None,
                        labels=["kind:package", "lang:js"],
                        attributes={"package": spec},
                    )
                self._edges.append(
                    Edge(
                        src=src_eid,
                        dst=eid,
                        type="imports",
                        attributes={"lang": "js_ts"},
                    )
                )

    def _resolve_local_js_like_spec(self, base_dir: str, spec: str) -> Optional[EntityId]:
        # Resolve "./foo" style specifiers to actual files if present.
        # Try with and without common JS/TS extensions.
        exts = ["", ".js", ".jsx", ".mjs", ".ts", ".tsx"]
        for ext in exts:
            cand = spec
            if not cand.endswith(ext):
                cand = spec + ext
            abs_cand = os.path.abspath(os.path.join(base_dir, cand))
            eid = self._file_entity_for_abs.get(abs_cand)
            if eid:
                return eid
        return None

    def _ingest_python_codegraph(self, cg: Any) -> None:
        """
        Integrate a scripts.codegraph_core.CodeGraph instance as one backend.

        This adds:
          - Module / function / class entities for Python.
          - "owns" edges from file entities to these symbols.
          - "imports" / "calls" / "tests" edges as provided by the backend.
        """
        # Entities
        for ent in cg.entities():
            abs_file = os.path.abspath(getattr(ent, "file", ""))
            file_eid = self._file_entity_for_abs.get(abs_file)
# ... truncated ...


from __future__ import annotations

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

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

def artifact_uri(program_id: str, rel_path: str) -> str:
    rel = rel_path.replace("\\", "/").lstrip("/")
    return f"program://{program_id}/artifact/{rel}"

def parse_program_uri(uri: str) -> Tuple[str, str, str, Optional[Tuple[int, int]]]:
    m = re.match(r"^program://([^/]+)/([^/]+)/(.+?)(?:#L(\d+)-L(\d+))?$", uri)
    if not m:
        raise ValueError(f"invalid program uri: {uri}")
    pid, kind, res, a, b = m.group(1), m.group(2), m.group(3), m.group(4), m.group(5)
    span = (int(a), int(b)) if (a and b) else None
    return pid, kind, res, span

class RepoGraph(ProgramGraph):
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.repo_root = os.path.abspath(repo_root)
        self.program_id = program_id_for_repo(self.repo_root)


from __future__ import annotations

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

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

def artifact_uri(program_id: str, rel_path: str) -> str:
    rel = rel_path.replace("\\", "/").lstrip("/")
    return f"program://{program_id}/artifact/{rel}"

def parse_program_uri(uri: str) -> Tuple[str, str, str, Optional[Tuple[int, int]]]:
    m = re.match(r"^program://([^/]+)/([^/]+)/(.+?)(?:#L(\d+)-L(\d+))?$", uri)
    if not m:
        raise ValueError(f"invalid program uri: {uri}")
    pid, kind, res, a, b = m.group(1), m.group(2), m.group(3), m.group(4), m.group(5)
    span = (int(a), int(b)) if (a and b) else None
    return pid, kind, res, span

class RepoGraph(ProgramGraph):
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.repo_root = os.path.abspath(repo_root)
        self.program_id = program_id_for_repo(self.repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self._file_hash: Dict[str, str] = {}
        # Cached graph views
        self._entities: Dict[EntityId, Entity] = {}
        self._edges: List[Edge] = []


from __future__ import annotations

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

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

def artifact_uri(program_id: str, rel_path: str) -> str:
    rel = rel_path.replace("\\", "/").lstrip("/")
    return f"program://{program_id}/artifact/{rel}"

def parse_program_uri(uri: str) -> Tuple[str, str, str, Optional[Tuple[int, int]]]:
    m = re.match(r"^program://([^/]+)/([^/]+)/(.+?)(?:#L(\d+)-L(\d+))?$", uri)
    if not m:
        raise ValueError(f"invalid program uri: {uri}")
    pid, kind, res, a, b = m.group(1), m.group(2), m.group(3), m.group(4), m.group(5)
    span = (int(a), int(b)) if (a and b) else None
    return pid, kind, res, span

class RepoGraph(ProgramGraph):
    def __init__(self, repo_root: str, ignore: Optional[List[str]] = None):
        self.repo_root = os.path.abspath(repo_root)
        self.program_id = program_id_for_repo(self.repo_root)
        self.ignore_rules = [s for s in (ignore or []) if s]
        self._file_hash: Dict[str, str] = {}
        # Cached graph views
        self._entities: Dict[EntityId, Entity] = {}
        self._edges: List[Edge] = []
        self._built: bool = False
        # Convenience indices
        #   - absolute file path → file-entity id
        #   - symbol name (lowercased) → list of entity ids
        self._file_entity_for_abs: Dict[str, EntityId] = {}
        self._symbols_by_name: Dict[str, List[EntityId]] = {}

    # ProgramGraph: core views
    def entities(self) -> Iterable[Entity]: