TOLBERTConfig

from dataclasses import dataclass
from typing import Any, Dict, Optional, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel

@dataclass
class TOLBERTConfig:
    """
    Minimal configuration for the TOLBERT encoder.

    This mirrors the sketch in `docs/api_reference.md` and is intended
    to be extended with any additional hyperparameters you need.
    """

    base_model_name: str
    level_sizes: Dict[int, int]
    proj_dim: int = 256
    # Overall weight on hierarchical classification loss
    lambda_hier: float = 1.0
    # Weight on path-consistency regularization term
    lambda_path: float = 0.0
    # Kept for completeness; used in the training script when adding contrastive loss.
    lambda_contrast: float = 0.0

class TOLBERT(nn.Module):
    """
    BERT/RoBERTa-style encoder with multi-level heads and a projection head.

    Forward returns:
      - loss: aggregated loss (MLM + hierarchical + optional path) if labels provided.
      - loss_components: dict with individual loss terms.
      - mlm_logits: token-level logits for MLM.
      - level_logits: dict[level_str] -> logits over nodes at that level.
      - proj: normalized CLS projection for contrastive / retrieval use.
    """

    def __init__(self, config: TOLBERTConfig):
        super().__init__()
        self.config = config

        # Backbone encoder
        self.encoder = AutoModel.from_pretrained(


TOLBERT

@dataclass
class TOLBERTConfig:
    """
    Minimal configuration for the TOLBERT encoder.

    This mirrors the sketch in `docs/api_reference.md` and is intended
    to be extended with any additional hyperparameters you need.
    """

    base_model_name: str
    level_sizes: Dict[int, int]
    proj_dim: int = 256
    # Overall weight on hierarchical classification loss
    lambda_hier: float = 1.0
    # Weight on path-consistency regularization term
    lambda_path: float = 0.0
    # Kept for completeness; used in the training script when adding contrastive loss.
    lambda_contrast: float = 0.0

class TOLBERT(nn.Module):
    """
    BERT/RoBERTa-style encoder with multi-level heads and a projection head.

    Forward returns:
      - loss: aggregated loss (MLM + hierarchical + optional path) if labels provided.
      - loss_components: dict with individual loss terms.
      - mlm_logits: token-level logits for MLM.
      - level_logits: dict[level_str] -> logits over nodes at that level.
      - proj: normalized CLS projection for contrastive / retrieval use.
    """

    def __init__(self, config: TOLBERTConfig):
        super().__init__()
        self.config = config

        # Backbone encoder
        self.encoder = AutoModel.from_pretrained(
            config.base_model_name,
            cache_dir="/data/checkpoints/",  # noqa: E501
        )
        hidden_dim = self.encoder.config.hidden_size

        # Simple MLM head (you can replace with the base model's own head)
        self.mlm_head = nn.Linear(hidden_dim, self.encoder.config.vocab_size)

        # Hierarchical heads per level
        self.level_heads = nn.ModuleDict()
        for level, size in config.level_sizes.items():
            self.level_heads[str(level)] = nn.Linear(hidden_dim, size)

        # Contrastive projection head
        self.proj = nn.Linear(hidden_dim, config.proj_dim)

        # Pre-sort levels once for stable ordering in losses
        self._sorted_levels = sorted(config.level_sizes.keys())

    @classmethod
    def from_pretrained(
        cls,
        checkpoint: str,
        config: Optional[TOLBERTConfig] = None,
    ) -> "TOLBERT":
        """
        Very minimal `from_pretrained` helper.

        - `checkpoint` is expected to be a `torch.save(model.state_dict())` file.
        - You must supply a `config` that matches the saved model.
        """
        if config is None:
            raise ValueError("TOLBERT.from_pretrained requires a TOLBERTConfig.")

        model = cls(config)
        state_dict = torch.load(checkpoint, map_location="cpu")
        model.load_state_dict(state_dict)
        return model

    def _compute_path_consistency_loss(
        self,
        level_logits: Dict[str, torch.Tensor],
        paths: Optional[Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]]],
    ) -> Optional[torch.Tensor]:
        """
        Batch-local implementation of the path-consistency loss described
        in the paper / training docs, extended to handle DAGs / multi-parent
        ontologies.

        For each adjacent level pair (ℓ-1, ℓ), we:
          - build a parent -> {child,...} mapping from the observed `paths`
            across the batch (allowing multiple parents per child),
          - compute softmax over level-ℓ logits,
          - roll level-ℓ probabilities up to the parent space and
            encourage agreement between:
                - the distribution implied by children, and
                - the model's own distribution at level ℓ-1.

        In a pure tree, each child has exactly one parent; in a DAG, a child
        may appear under multiple parents across the provided paths. This
        implementation treats all observed parent-child relations as valid.
        """
        if paths is None:
            return None

        if len(self._sorted_levels) < 2:
            return None

        device = next(self.parameters()).device

        # Normalize paths into per-example path sets to support DAGs.
        # Each element in `path_sets` is a list of paths (lists of node ids).
        path_sets: list[list[list[int]]] = []
        for p in paths:
            if not p:
                path_sets.append([])
                continue
            first = p[0]
            if isinstance(first, int):
                path_sets.append([[int(x) for x in p]])  # type: ignore[arg-type]
            else:
                path_set: list[list[int]] = []
                for sub in p:  # type: ignore[assignment]
                    if isinstance(sub, (list, tuple)):
                        path_set.append([int(x) for x in sub])
                path_sets.append(path_set)

        batch_size = len(path_sets)

        total_kl = 0.0
        num_terms = 0

        # We assume level index ℓ in the paper corresponds to path index ℓ,
        # with ℓ = 1..L (and index 0 being the root, which we do not have a head for).
        # For model heads we use integer "levels" taken from config.level_sizes.
        for idx in range(1, len(self._sorted_levels)):
            level_prev = self._sorted_levels[idx - 1]
            level_curr = self._sorted_levels[idx]

            if str(level_prev) not in level_logits or str(level_curr) not in level_logits:
                continue

            logits_prev = level_logits[str(level_prev)]  # (B, C_prev)
            logits_curr = level_logits[str(level_curr)]  # (B, C_curr)

            probs_prev = F.softmax(logits_prev, dim=-1)  # (B, C_prev)
            probs_curr = F.softmax(logits_curr, dim=-1)  # (B, C_curr)

            # Build parent -> children map for this level-pair from batch paths.
            parent_to_children: Dict[int, set[int]] = {}
            for b in range(batch_size):
                for path in path_sets[b]:
                    if len(path) <= max(level_prev, level_curr):
                        continue
                    parent_id = path[level_prev]
                    child_id = path[level_curr]
                    if parent_id < 0 or child_id < 0:
                        continue
                    if parent_id not in parent_to_children:
                        parent_to_children[parent_id] = set()
                    parent_to_children[parent_id].add(child_id)

            if not parent_to_children:
                continue

            num_parents = logits_prev.size(-1)
            rolled = torch.zeros_like(probs_prev)
            for parent_id, children_ids in parent_to_children.items():
                if parent_id < 0 or parent_id >= num_parents:
                    continue
                children_idx = torch.tensor(
                    list(children_ids),
                    dtype=torch.long,
                    device=device,
                )
                children_idx = children_idx[
                    (children_idx >= 0) & (children_idx < probs_curr.size(-1))
                ]
                if children_idx.numel() == 0:
                    continue
                rolled[:, parent_id] = probs_curr[:, children_idx].sum(dim=-1)

            # Renormalize rolled-up distribution to avoid degenerate zeros.
            rolled_sum = rolled.sum(dim=-1, keepdim=True)
            rolled_sum = torch.clamp(rolled_sum, min=1e-8)
            rolled = rolled / rolled_sum

            # Ensure numerical stability for KL
            probs_prev_clamped = torch.clamp(probs_prev, min=1e-8)
            rolled_clamped = torch.clamp(rolled, min=1e-8)

            kl = F.kl_div(
                probs_prev_clamped.log(),
                rolled_clamped,
                reduction="batchmean",
            )
            total_kl = total_kl + kl
            num_terms += 1

        if num_terms == 0:
            return None

        return total_kl / float(num_terms)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels_mlm: Optional[torch.Tensor] = None,
        level_targets: Optional[Dict[int, torch.Tensor]] = None,
        paths: Optional[Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]]] = None,
    ) -> Dict[str, Any]:
        """
        Args:
            input_ids: (batch, seq_len)
            attention_mask: (batch, seq_len)
            labels_mlm: (batch, seq_len) with -100 for non-MLM positions.
            level_targets: dict {level(int): tensor(batch,)} with class indices.
            paths: optional sequence of per-example label paths, each a
                   sequence of node ids [root, c1, c2, ..., cL]. Used for
                   path-consistency loss.
        """
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = outputs.last_hidden_state  # (batch, seq, hidden)
        cls = hidden_states[:, 0, :]  # (batch, hidden)

        # MLM logits
        mlm_logits = self.mlm_head(hidden_states)  # (batch, seq, vocab)

        # Hierarchical logits
        level_logits: Dict[str, torch.Tensor] = {}
        for level, head in self.level_heads.items():
            level_logits[level] = head(cls)  # (batch, C_level)

        # Contrastive projection
        proj = F.normalize(self.proj(cls), dim=-1)

        loss: Optional[torch.Tensor] = None
        loss_dict: Dict[str, torch.Tensor] = {}

        # MLM loss
        mlm_loss = None
        if labels_mlm is not None:
            mlm_loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            mlm_loss = mlm_loss_fct(
                mlm_logits.view(-1, mlm_logits.size(-1)),
                labels_mlm.view(-1),
            )
            loss_dict["mlm"] = mlm_loss

        # Hierarchical classification loss (sum over levels, excluding ignore_index)
        hier_loss = None
        if level_targets is not None and len(level_targets) > 0:
            ce = nn.CrossEntropyLoss(ignore_index=-100)
            per_level_losses = []
            for level_int in self._sorted_levels:
                if level_int not in level_targets:
                    continue
                logits = level_logits[str(level_int)]
                targets = level_targets[level_int]
                level_loss = ce(logits, targets)
                loss_dict[f"level_{level_int}"] = level_loss
                per_level_losses.append(level_loss)

            if per_level_losses:
                hier_loss = torch.stack(per_level_losses).mean()
                loss_dict["hier"] = hier_loss

        # Path-consistency loss (uses predicted distributions only)
        path_loss = None
        if paths is not None:
            path_loss = self._compute_path_consistency_loss(level_logits, paths)
            if path_loss is not None:
                loss_dict["path"] = path_loss

        # Aggregate total loss with configuration weights.
        components = []
        if mlm_loss is not None:
            components.append(mlm_loss)
        if hier_loss is not None and self.config.lambda_hier != 0.0:
            components.append(self.config.lambda_hier * hier_loss)
        if path_loss is not None and self.config.lambda_path != 0.0:
            components.append(self.config.lambda_path * path_loss)

        if components:
            loss = sum(components)

        return {
            "loss": loss,
            "loss_components": loss_dict,
            "mlm_logits": mlm_logits,
            "level_logits": level_logits,
            "proj": proj,
        }


    # Overall weight on hierarchical classification loss
    lambda_hier: float = 1.0
    # Weight on path-consistency regularization term
    lambda_path: float = 0.0
    # Kept for completeness; used in the training script when adding contrastive loss.
    lambda_contrast: float = 0.0

class TOLBERT(nn.Module):
    """
    BERT/RoBERTa-style encoder with multi-level heads and a projection head.

    Forward returns:
      - loss: aggregated loss (MLM + hierarchical + optional path) if labels provided.
      - loss_components: dict with individual loss terms.
      - mlm_logits: token-level logits for MLM.
      - level_logits: dict[level_str] -> logits over nodes at that level.
      - proj: normalized CLS projection for contrastive / retrieval use.
    """

    def __init__(self, config: TOLBERTConfig):
        super().__init__()
        self.config = config

        # Backbone encoder
        self.encoder = AutoModel.from_pretrained(
            config.base_model_name,
            cache_dir="/data/checkpoints/",  # noqa: E501
        )
        hidden_dim = self.encoder.config.hidden_size

        # Simple MLM head (you can replace with the base model's own head)
        self.mlm_head = nn.Linear(hidden_dim, self.encoder.config.vocab_size)

        # Hierarchical heads per level
        self.level_heads = nn.ModuleDict()
        for level, size in config.level_sizes.items():
            self.level_heads[str(level)] = nn.Linear(hidden_dim, size)

        # Contrastive projection head
        self.proj = nn.Linear(hidden_dim, config.proj_dim)

        # Pre-sort levels once for stable ordering in losses
        self._sorted_levels = sorted(config.level_sizes.keys())

    @classmethod
    def from_pretrained(
        cls,
        checkpoint: str,
        config: Optional[TOLBERTConfig] = None,
    ) -> "TOLBERT":
        """
        Very minimal `from_pretrained` helper.

        - `checkpoint` is expected to be a `torch.save(model.state_dict())` file.
        - You must supply a `config` that matches the saved model.
        """
        if config is None:
            raise ValueError("TOLBERT.from_pretrained requires a TOLBERTConfig.")

        model = cls(config)
        state_dict = torch.load(checkpoint, map_location="cpu")
        model.load_state_dict(state_dict)
        return model


            config.base_model_name,
            cache_dir="/data/checkpoints/",  # noqa: E501
        )
        hidden_dim = self.encoder.config.hidden_size

        # Simple MLM head (you can replace with the base model's own head)
        self.mlm_head = nn.Linear(hidden_dim, self.encoder.config.vocab_size)

        # Hierarchical heads per level
        self.level_heads = nn.ModuleDict()
        for level, size in config.level_sizes.items():
            self.level_heads[str(level)] = nn.Linear(hidden_dim, size)

        # Contrastive projection head
        self.proj = nn.Linear(hidden_dim, config.proj_dim)

        # Pre-sort levels once for stable ordering in losses
        self._sorted_levels = sorted(config.level_sizes.keys())

    @classmethod
    def from_pretrained(
        cls,
        checkpoint: str,
        config: Optional[TOLBERTConfig] = None,
    ) -> "TOLBERT":
        """
        Very minimal `from_pretrained` helper.

        - `checkpoint` is expected to be a `torch.save(model.state_dict())` file.
        - You must supply a `config` that matches the saved model.
        """
        if config is None:
            raise ValueError("TOLBERT.from_pretrained requires a TOLBERTConfig.")

        model = cls(config)
        state_dict = torch.load(checkpoint, map_location="cpu")
        model.load_state_dict(state_dict)
        return model

    def _compute_path_consistency_loss(
        self,
        level_logits: Dict[str, torch.Tensor],
        paths: Optional[Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]]],
    ) -> Optional[torch.Tensor]:
        """
        Batch-local implementation of the path-consistency loss described
        in the paper / training docs, extended to handle DAGs / multi-parent
        ontologies.

        For each adjacent level pair (ℓ-1, ℓ), we:
          - build a parent -> {child,...} mapping from the observed `paths`
            across the batch (allowing multiple parents per child),
          - compute softmax over level-ℓ logits,
          - roll level-ℓ probabilities up to the parent space and
            encourage agreement between:
                - the distribution implied by children, and
                - the model's own distribution at level ℓ-1.


    @classmethod
    def from_pretrained(
        cls,
        checkpoint: str,
        config: Optional[TOLBERTConfig] = None,
    ) -> "TOLBERT":
        """
        Very minimal `from_pretrained` helper.

        - `checkpoint` is expected to be a `torch.save(model.state_dict())` file.
        - You must supply a `config` that matches the saved model.
        """
        if config is None:
            raise ValueError("TOLBERT.from_pretrained requires a TOLBERTConfig.")

        model = cls(config)
        state_dict = torch.load(checkpoint, map_location="cpu")
        model.load_state_dict(state_dict)
        return model

    def _compute_path_consistency_loss(
        self,
        level_logits: Dict[str, torch.Tensor],
        paths: Optional[Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]]],
    ) -> Optional[torch.Tensor]:
        """
        Batch-local implementation of the path-consistency loss described
        in the paper / training docs, extended to handle DAGs / multi-parent
        ontologies.

        For each adjacent level pair (ℓ-1, ℓ), we:
          - build a parent -> {child,...} mapping from the observed `paths`
            across the batch (allowing multiple parents per child),
          - compute softmax over level-ℓ logits,
          - roll level-ℓ probabilities up to the parent space and
            encourage agreement between:
                - the distribution implied by children, and
                - the model's own distribution at level ℓ-1.

        In a pure tree, each child has exactly one parent; in a DAG, a child
        may appear under multiple parents across the provided paths. This
        implementation treats all observed parent-child relations as valid.
        """
        if paths is None:
            return None

        if len(self._sorted_levels) < 2:
            return None

        device = next(self.parameters()).device

        # Normalize paths into per-example path sets to support DAGs.
        # Each element in `path_sets` is a list of paths (lists of node ids).
        path_sets: list[list[list[int]]] = []
        for p in paths:
            if not p:
                path_sets.append([])
                continue
            first = p[0]
            if isinstance(first, int):
                path_sets.append([[int(x) for x in p]])  # type: ignore[arg-type]
            else:
                path_set: list[list[int]] = []
                for sub in p:  # type: ignore[assignment]
                    if isinstance(sub, (list, tuple)):
                        path_set.append([int(x) for x in sub])
                path_sets.append(path_set)

        batch_size = len(path_sets)

        total_kl = 0.0
        num_terms = 0

        # We assume level index ℓ in the paper corresponds to path index ℓ,
        # with ℓ = 1..L (and index 0 being the root, which we do not have a head for).
        # For model heads we use integer "levels" taken from config.level_sizes.
        for idx in range(1, len(self._sorted_levels)):
            level_prev = self._sorted_levels[idx - 1]
            level_curr = self._sorted_levels[idx]

            if str(level_prev) not in level_logits or str(level_curr) not in level_logits:
                continue

            logits_prev = level_logits[str(level_prev)]  # (B, C_prev)
            logits_curr = level_logits[str(level_curr)]  # (B, C_curr)

            probs_prev = F.softmax(logits_prev, dim=-1)  # (B, C_prev)
            probs_curr = F.softmax(logits_curr, dim=-1)  # (B, C_curr)

            # Build parent -> children map for this level-pair from batch paths.
            parent_to_children: Dict[int, set[int]] = {}
            for b in range(batch_size):
                for path in path_sets[b]:
                    if len(path) <= max(level_prev, level_curr):
                        continue
                    parent_id = path[level_prev]
                    child_id = path[level_curr]
                    if parent_id < 0 or child_id < 0:
                        continue
                    if parent_id not in parent_to_children:
                        parent_to_children[parent_id] = set()
                    parent_to_children[parent_id].add(child_id)

            if not parent_to_children:
                continue

            num_parents = logits_prev.size(-1)
            rolled = torch.zeros_like(probs_prev)
            for parent_id, children_ids in parent_to_children.items():
                if parent_id < 0 or parent_id >= num_parents:
                    continue
                children_idx = torch.tensor(
                    list(children_ids),
                    dtype=torch.long,
                    device=device,
                )
                children_idx = children_idx[
                    (children_idx >= 0) & (children_idx < probs_curr.size(-1))
                ]
                if children_idx.numel() == 0:
                    continue
                rolled[:, parent_id] = probs_curr[:, children_idx].sum(dim=-1)

            # Renormalize rolled-up distribution to avoid degenerate zeros.
            rolled_sum = rolled.sum(dim=-1, keepdim=True)
            rolled_sum = torch.clamp(rolled_sum, min=1e-8)
            rolled = rolled / rolled_sum

            # Ensure numerical stability for KL
            probs_prev_clamped = torch.clamp(probs_prev, min=1e-8)
            rolled_clamped = torch.clamp(rolled, min=1e-8)

            kl = F.kl_div(
                probs_prev_clamped.log(),
                rolled_clamped,
                reduction="batchmean",
            )
            total_kl = total_kl + kl
            num_terms += 1

        if num_terms == 0:
            return None

        return total_kl / float(num_terms)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels_mlm: Optional[torch.Tensor] = None,
        level_targets: Optional[Dict[int, torch.Tensor]] = None,
        paths: Optional[Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]]] = None,
    ) -> Dict[str, Any]:
        """
        Args:
            input_ids: (batch, seq_len)
            attention_mask: (batch, seq_len)
            labels_mlm: (batch, seq_len) with -100 for non-MLM positions.
            level_targets: dict {level(int): tensor(batch,)} with class indices.
            paths: optional sequence of per-example label paths, each a
                   sequence of node ids [root, c1, c2, ..., cL]. Used for
                   path-consistency loss.
        """
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)


forward

            rolled_sum = torch.clamp(rolled_sum, min=1e-8)
            rolled = rolled / rolled_sum

            # Ensure numerical stability for KL
            probs_prev_clamped = torch.clamp(probs_prev, min=1e-8)
            rolled_clamped = torch.clamp(rolled, min=1e-8)

            kl = F.kl_div(
                probs_prev_clamped.log(),
                rolled_clamped,
                reduction="batchmean",
            )
            total_kl = total_kl + kl
            num_terms += 1

        if num_terms == 0:
            return None

        return total_kl / float(num_terms)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels_mlm: Optional[torch.Tensor] = None,
        level_targets: Optional[Dict[int, torch.Tensor]] = None,
        paths: Optional[Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]]] = None,
    ) -> Dict[str, Any]:
        """
        Args:
            input_ids: (batch, seq_len)
            attention_mask: (batch, seq_len)
            labels_mlm: (batch, seq_len) with -100 for non-MLM positions.
            level_targets: dict {level(int): tensor(batch,)} with class indices.
            paths: optional sequence of per-example label paths, each a
                   sequence of node ids [root, c1, c2, ..., cL]. Used for
                   path-consistency loss.
        """
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = outputs.last_hidden_state  # (batch, seq, hidden)
        cls = hidden_states[:, 0, :]  # (batch, hidden)

        # MLM logits
        mlm_logits = self.mlm_head(hidden_states)  # (batch, seq, vocab)

        # Hierarchical logits
        level_logits: Dict[str, torch.Tensor] = {}
        for level, head in self.level_heads.items():
            level_logits[level] = head(cls)  # (batch, C_level)

        # Contrastive projection
        proj = F.normalize(self.proj(cls), dim=-1)

        loss: Optional[torch.Tensor] = None
        loss_dict: Dict[str, torch.Tensor] = {}

        # MLM loss
        mlm_loss = None
        if labels_mlm is not None:
            mlm_loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            mlm_loss = mlm_loss_fct(
                mlm_logits.view(-1, mlm_logits.size(-1)),
                labels_mlm.view(-1),
            )
            loss_dict["mlm"] = mlm_loss

        # Hierarchical classification loss (sum over levels, excluding ignore_index)
        hier_loss = None
        if level_targets is not None and len(level_targets) > 0:
            ce = nn.CrossEntropyLoss(ignore_index=-100)
            per_level_losses = []
            for level_int in self._sorted_levels:
                if level_int not in level_targets:
                    continue
                logits = level_logits[str(level_int)]
                targets = level_targets[level_int]
                level_loss = ce(logits, targets)
                loss_dict[f"level_{level_int}"] = level_loss
                per_level_losses.append(level_loss)

            if per_level_losses:
                hier_loss = torch.stack(per_level_losses).mean()
                loss_dict["hier"] = hier_loss

        # Path-consistency loss (uses predicted distributions only)
        path_loss = None
        if paths is not None:
            path_loss = self._compute_path_consistency_loss(level_logits, paths)
            if path_loss is not None:
                loss_dict["path"] = path_loss

        # Aggregate total loss with configuration weights.
        components = []
        if mlm_loss is not None:
            components.append(mlm_loss)
        if hier_loss is not None and self.config.lambda_hier != 0.0:
            components.append(self.config.lambda_hier * hier_loss)
        if path_loss is not None and self.config.lambda_path != 0.0:
            components.append(self.config.lambda_path * path_loss)

        if components:
            loss = sum(components)

        return {
            "loss": loss,
            "loss_components": loss_dict,
            "mlm_logits": mlm_logits,
            "level_logits": level_logits,
            "proj": proj,
        }


tolbert.losses

from typing import List, Sequence

import torch

def _compute_shared_depth(pi: Sequence[int], pj: Sequence[int]) -> int:
    """
    Depth of deepest shared label beyond the root.

    We treat index 0 in the path as the root. The shared depth is the
    number of consecutive matching labels starting from index 1.
    """
    max_len = min(len(pi), len(pj))
    depth = 0
    # Start from index 1 to skip the root.
    for idx in range(1, max_len):
        if pi[idx] != pj[idx]:
            break
        depth += 1
    return depth

def _normalize_paths_arg(
    paths: Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]],
) -> List[List[List[int]]]:
    """
    Normalize `paths` into a list of per-example path sets.

    Accepted input formats:
      - Tree-style (single path per example):
            paths = [[v0, v1, ...], [v0, v1', ...], ...]
      - DAG-style (multiple valid paths per example):
            paths = [
              [[v0, v1, ...], [v0, v1', ...]],
              [[v0, u1, ...]],
              ...
            ]

    Output format:
        List[ List[List[int]] ], where outer index is batch example and each
        inner list is one valid path for that example.
    """
    norm: List[List[List[int]]] = []
    for p in paths:
        if not p:
            norm.append([])
            continue
        first = p[0]
        # If first element is an int, we have a single path.
        if isinstance(first, int):
            norm.append([[int(x) for x in p]])  # type: ignore[arg-type]
        else:
            # Assume list/tuple of paths.
            path_set: List[List[int]] = []
            for sub in p:  # type: ignore[assignment]
                if isinstance(sub, (list, tuple)):
                    path_set.append([int(x) for x in sub])
            norm.append(path_set)
    return norm

def _compute_shared_depth_sets(
    paths_i: List[List[int]],
    paths_j: List[List[int]],
) -> int:
    """
    Deepest shared depth between two *sets* of paths (for DAGs).

    We take the maximum shared depth over all path pairs (pi in paths_i,
    pj in paths_j).
    """
    max_depth = 0
    for pi in paths_i:
        for pj in paths_j:
            d = _compute_shared_depth(pi, pj)
            if d > max_depth:
                max_depth = d
    return max_depth

def tree_contrastive_loss(
    embeddings: torch.Tensor,
    paths: Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]],
    temperature: float = 0.07,
) -> torch.Tensor:
    """
    Depth-weighted tree-aware InfoNCE-style contrastive loss.

    This implements the hierarchical supervised contrastive objective
    described in the paper, extended to handle DAGs / multiple paths per
    example:

      - Each example i has one or more valid paths
            (y_1^{(i,m)}, ..., y_L^{(i,m)})
        through the ontology graph.
      - For a pair (i, j), define depth(i, j) as the deepest level at
        which they share the same label beyond the root **along any**
        of their respective paths.
      - Positives P(i) are all j != i with depth(i, j) >= 1.
      - Each positive pair is weighted by w_ij ∝ depth(i, j).
      - Negatives are examples with depth(i, j) == 0 (only root shared).

    Args:
        embeddings: (batch, dim) L2-normalized embeddings (e.g., proj outputs).
        paths: sequence of per-example paths, which may be:
            - a single path per example: [[v0, v1, ...], ...], or
            - multiple paths per example: [[[v0, v1, ...], [v0, v1', ...]], ...]
        temperature: softmax temperature for contrastive logits.
    """
    device = embeddings.device
    batch_size = embeddings.size(0)

    if batch_size <= 1:
        return torch.tensor(0.0, device=device)

    # Normalize paths into per-example path sets to support DAGs.
    path_sets = _normalize_paths_arg(paths)

    # Compute pairwise shared depths over path sets.
    depth_mat = torch.zeros((batch_size, batch_size), dtype=torch.float32, device=device)
    for i in range(batch_size):
        for j in range(batch_size):
            if i == j:
                continue
            d = _compute_shared_depth_sets(path_sets[i], path_sets[j])
            depth_mat[i, j] = float(d)

    # Positives: depth >= 1
    pos_mask = depth_mat >= 1.0  # (B, B)

    # If no pairs share any non-root ancestor, the contrastive signal is empty.
    if not pos_mask.any():
        return torch.tensor(0.0, device=device)

    # Similarity matrix
    sim = embeddings @ embeddings.T  # (B, B)
    sim = sim / temperature

    # Exclude self from denominators by setting diag to a very large
    # negative value (rather than -inf) to avoid NaNs when multiplied
    # by zero weights later on.
    diag_mask = torch.eye(batch_size, dtype=torch.bool, device=device)
    sim = sim.masked_fill(diag_mask, -1e9)

    # Log-softmax-like normalization per row
    logsumexp = torch.logsumexp(sim, dim=1, keepdim=True)  # (B, 1)
    log_probs = sim - logsumexp  # (B, B)

    # Depth-based weights w_ij = depth(i,j) / L_i for positives, where
    # L_i is the maximum usable depth for anchor i (excluding root).
    # For DAGs, we use the maximum depth over that example's path set.
    anchor_depths: List[int] = []
    for path_set in path_sets:
        max_len = 1
        for p in path_set:
            if len(p) - 1 > max_len:
                max_len = len(p) - 1
        anchor_depths.append(max(1, max_len))
    L = torch.tensor(anchor_depths, dtype=torch.float32, device=device).unsqueeze(1)  # (B, 1)

    # Avoid divide-by-zero; we clamped to at least 1 above.
    norm_depths = depth_mat / L

    weights = torch.zeros_like(depth_mat)
    weights[pos_mask] = norm_depths[pos_mask]

    # For each anchor i, compute a weighted average of log_probs over positives.
    eps = 1e-8
    row_pos_counts = pos_mask.sum(dim=1)  # (B,)
    row_weight_sums = weights.sum(dim=1)  # (B,)

    # Avoid divide-by-zero
    valid_rows = row_pos_counts > 0
    if not valid_rows.any():
        return torch.tensor(0.0, device=device)

    # Normalized weights within each row (only over positives)
    norm_weights = torch.zeros_like(weights)
    norm_weights[pos_mask] = (
        weights[pos_mask]
        / (row_weight_sums.unsqueeze(1).expand_as(weights)[pos_mask] + eps)
    )

    # For each row i, loss_i = - (1/|P(i)|) Σ_j w_ij_norm * log_probs_ij
    weighted_log_probs = (norm_weights * log_probs)  # (B, B)
    row_losses = torch.zeros(batch_size, device=device)

    # Sum over j for each i
    row_losses = -weighted_log_probs.sum(dim=1)
    # Divide by |P(i)| for anchors that have positives
    row_losses = torch.where(
        valid_rows,
        row_losses / (row_pos_counts.to(row_losses.dtype) + eps),
        torch.zeros_like(row_losses),
    )

    # Average over anchors with at least one positive
    loss = row_losses[valid_rows].mean()
    return loss


from typing import List, Sequence

import torch

def _compute_shared_depth(pi: Sequence[int], pj: Sequence[int]) -> int:
    """
    Depth of deepest shared label beyond the root.

    We treat index 0 in the path as the root. The shared depth is the
    number of consecutive matching labels starting from index 1.
    """
    max_len = min(len(pi), len(pj))
    depth = 0
    # Start from index 1 to skip the root.
    for idx in range(1, max_len):
        if pi[idx] != pj[idx]:
            break
        depth += 1
    return depth

def _normalize_paths_arg(
    paths: Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]],
) -> List[List[List[int]]]:
    """
    Normalize `paths` into a list of per-example path sets.

    Accepted input formats:
      - Tree-style (single path per example):
            paths = [[v0, v1, ...], [v0, v1', ...], ...]
      - DAG-style (multiple valid paths per example):
            paths = [
              [[v0, v1, ...], [v0, v1', ...]],
              [[v0, u1, ...]],
              ...
            ]

    Output format:
        List[ List[List[int]] ], where outer index is batch example and each


import torch

def _compute_shared_depth(pi: Sequence[int], pj: Sequence[int]) -> int:
    """
    Depth of deepest shared label beyond the root.

    We treat index 0 in the path as the root. The shared depth is the
    number of consecutive matching labels starting from index 1.
    """
    max_len = min(len(pi), len(pj))
    depth = 0
    # Start from index 1 to skip the root.
    for idx in range(1, max_len):
        if pi[idx] != pj[idx]:
            break
        depth += 1
    return depth

def _normalize_paths_arg(
    paths: Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]],
) -> List[List[List[int]]]:
    """
    Normalize `paths` into a list of per-example path sets.

    Accepted input formats:
      - Tree-style (single path per example):
            paths = [[v0, v1, ...], [v0, v1', ...], ...]
      - DAG-style (multiple valid paths per example):
            paths = [
              [[v0, v1, ...], [v0, v1', ...]],
              [[v0, u1, ...]],
              ...
            ]

    Output format:
        List[ List[List[int]] ], where outer index is batch example and each
        inner list is one valid path for that example.
    """
    norm: List[List[List[int]]] = []
    for p in paths:
        if not p:
            norm.append([])
            continue
        first = p[0]
        # If first element is an int, we have a single path.
        if isinstance(first, int):
            norm.append([[int(x) for x in p]])  # type: ignore[arg-type]
        else:
            # Assume list/tuple of paths.
            path_set: List[List[int]] = []
            for sub in p:  # type: ignore[assignment]
                if isinstance(sub, (list, tuple)):
                    path_set.append([int(x) for x in sub])
            norm.append(path_set)
    return norm

def _compute_shared_depth_sets(
    paths_i: List[List[int]],
    paths_j: List[List[int]],
) -> int:
    """
    Deepest shared depth between two *sets* of paths (for DAGs).

    We take the maximum shared depth over all path pairs (pi in paths_i,
    pj in paths_j).
    """
    max_depth = 0
    for pi in paths_i:
        for pj in paths_j:
            d = _compute_shared_depth(pi, pj)
            if d > max_depth:
                max_depth = d
    return max_depth


    """
    norm: List[List[List[int]]] = []
    for p in paths:
        if not p:
            norm.append([])
            continue
        first = p[0]
        # If first element is an int, we have a single path.
        if isinstance(first, int):
            norm.append([[int(x) for x in p]])  # type: ignore[arg-type]
        else:
            # Assume list/tuple of paths.
            path_set: List[List[int]] = []
            for sub in p:  # type: ignore[assignment]
                if isinstance(sub, (list, tuple)):
                    path_set.append([int(x) for x in sub])
            norm.append(path_set)
    return norm

def _compute_shared_depth_sets(
    paths_i: List[List[int]],
    paths_j: List[List[int]],
) -> int:
    """
    Deepest shared depth between two *sets* of paths (for DAGs).

    We take the maximum shared depth over all path pairs (pi in paths_i,
    pj in paths_j).
    """
    max_depth = 0
    for pi in paths_i:
        for pj in paths_j:
            d = _compute_shared_depth(pi, pj)
            if d > max_depth:
                max_depth = d
    return max_depth

def tree_contrastive_loss(
    embeddings: torch.Tensor,
    paths: Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]],
    temperature: float = 0.07,
) -> torch.Tensor:
    """
    Depth-weighted tree-aware InfoNCE-style contrastive loss.

    This implements the hierarchical supervised contrastive objective
    described in the paper, extended to handle DAGs / multiple paths per
    example:

      - Each example i has one or more valid paths
            (y_1^{(i,m)}, ..., y_L^{(i,m)})
        through the ontology graph.
      - For a pair (i, j), define depth(i, j) as the deepest level at
        which they share the same label beyond the root **along any**
        of their respective paths.


def _compute_shared_depth_sets(
    paths_i: List[List[int]],
    paths_j: List[List[int]],
) -> int:
    """
    Deepest shared depth between two *sets* of paths (for DAGs).

    We take the maximum shared depth over all path pairs (pi in paths_i,
    pj in paths_j).
    """
    max_depth = 0
    for pi in paths_i:
        for pj in paths_j:
            d = _compute_shared_depth(pi, pj)
            if d > max_depth:
                max_depth = d
    return max_depth

def tree_contrastive_loss(
    embeddings: torch.Tensor,
    paths: Sequence[Sequence[int]] | Sequence[Sequence[Sequence[int]]],
    temperature: float = 0.07,
) -> torch.Tensor:
    """
    Depth-weighted tree-aware InfoNCE-style contrastive loss.

    This implements the hierarchical supervised contrastive objective
    described in the paper, extended to handle DAGs / multiple paths per
    example:

      - Each example i has one or more valid paths
            (y_1^{(i,m)}, ..., y_L^{(i,m)})
        through the ontology graph.
      - For a pair (i, j), define depth(i, j) as the deepest level at
        which they share the same label beyond the root **along any**
        of their respective paths.
      - Positives P(i) are all j != i with depth(i, j) >= 1.
      - Each positive pair is weighted by w_ij ∝ depth(i, j).
      - Negatives are examples with depth(i, j) == 0 (only root shared).

    Args:
        embeddings: (batch, dim) L2-normalized embeddings (e.g., proj outputs).
        paths: sequence of per-example paths, which may be:
            - a single path per example: [[v0, v1, ...], ...], or
            - multiple paths per example: [[[v0, v1, ...], [v0, v1', ...]], ...]
        temperature: softmax temperature for contrastive logits.
    """
    device = embeddings.device
    batch_size = embeddings.size(0)

    if batch_size <= 1:
        return torch.tensor(0.0, device=device)

    # Normalize paths into per-example path sets to support DAGs.
    path_sets = _normalize_paths_arg(paths)

    # Compute pairwise shared depths over path sets.
    depth_mat = torch.zeros((batch_size, batch_size), dtype=torch.float32, device=device)
    for i in range(batch_size):
        for j in range(batch_size):
            if i == j:
                continue
            d = _compute_shared_depth_sets(path_sets[i], path_sets[j])
            depth_mat[i, j] = float(d)

    # Positives: depth >= 1
    pos_mask = depth_mat >= 1.0  # (B, B)

    # If no pairs share any non-root ancestor, the contrastive signal is empty.
    if not pos_mask.any():
        return torch.tensor(0.0, device=device)

    # Similarity matrix
    sim = embeddings @ embeddings.T  # (B, B)
    sim = sim / temperature

    # Exclude self from denominators by setting diag to a very large
    # negative value (rather than -inf) to avoid NaNs when multiplied
    # by zero weights later on.
    diag_mask = torch.eye(batch_size, dtype=torch.bool, device=device)
    sim = sim.masked_fill(diag_mask, -1e9)

    # Log-softmax-like normalization per row
    logsumexp = torch.logsumexp(sim, dim=1, keepdim=True)  # (B, 1)
    log_probs = sim - logsumexp  # (B, B)

    # Depth-based weights w_ij = depth(i,j) / L_i for positives, where
    # L_i is the maximum usable depth for anchor i (excluding root).
    # For DAGs, we use the maximum depth over that example's path set.
    anchor_depths: List[int] = []
    for path_set in path_sets:
        max_len = 1
        for p in path_set:
            if len(p) - 1 > max_len:
                max_len = len(p) - 1
        anchor_depths.append(max(1, max_len))
    L = torch.tensor(anchor_depths, dtype=torch.float32, device=device).unsqueeze(1)  # (B, 1)

    # Avoid divide-by-zero; we clamped to at least 1 above.
    norm_depths = depth_mat / L

    weights = torch.zeros_like(depth_mat)
    weights[pos_mask] = norm_depths[pos_mask]

    # For each anchor i, compute a weighted average of log_probs over positives.
    eps = 1e-8
    row_pos_counts = pos_mask.sum(dim=1)  # (B,)
    row_weight_sums = weights.sum(dim=1)  # (B,)

    # Avoid divide-by-zero
    valid_rows = row_pos_counts > 0
    if not valid_rows.any():
        return torch.tensor(0.0, device=device)

    # Normalized weights within each row (only over positives)
    norm_weights = torch.zeros_like(weights)
    norm_weights[pos_mask] = (
        weights[pos_mask]
        / (row_weight_sums.unsqueeze(1).expand_as(weights)[pos_mask] + eps)
    )

    # For each row i, loss_i = - (1/|P(i)|) Σ_j w_ij_norm * log_probs_ij
    weighted_log_probs = (norm_weights * log_probs)  # (B, B)
    row_losses = torch.zeros(batch_size, device=device)

    # Sum over j for each i
    row_losses = -weighted_log_probs.sum(dim=1)
    # Divide by |P(i)| for anchors that have positives
    row_losses = torch.where(
        valid_rows,
        row_losses / (row_pos_counts.to(row_losses.dtype) + eps),
        torch.zeros_like(row_losses),
    )

    # Average over anchors with at least one positive
    loss = row_losses[valid_rows].mean()
    return loss


tolbert.data

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Any

import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizerBase

@dataclass
class SpanRecord:
    text: str
    # Canonical single path used for per-level classification targets.
    node_path: Optional[List[int]]
    # Full raw record, which may include richer DAG-style fields such as
    # `node_paths` (multiple valid paths) in addition to `node_path`.
    raw: Dict[str, Any]

class TreeOfLifeDataset(Dataset):
    """
    Minimal dataset for TOLBERT training.

    Assumes a JSONL `spans_file` where each line looks like:
        {
          "span_id": "s0",
          "text": "def forward(...):",
          "node_path": [0, 1, 2, 10]   # optional but recommended
        }

    You are free to extend this to include per-level fields instead of
    `node_path` (e.g., `level_1_id`, `level_2_id`, ...). This class is
    intentionally simple and meant as a starting point.
    """

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

        self._records: List[SpanRecord] = []
        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)
                text = obj["text"]

                # Support both legacy single-path and richer DAG-style formats:
                #   - "node_path": [v0, v1, ..., vK]
                #   - "node_paths": [[v0, v1, ..., vK], [v0, v1', ..., vK'], ...]
                #
                # We keep a single canonical path (first path if multiple are
                # provided) for per-level classification, while preserving the
                # richer structure in `raw` for path-consistency and
                # contrastive losses.
                node_path: Optional[List[int]] = None
                if "node_paths" in obj and obj["node_paths"] is not None:
                    paths_val = obj["node_paths"]
                    if isinstance(paths_val, list) and paths_val:
                        first = paths_val[0]
                        if isinstance(first, (list, tuple)):
                            node_path = [int(x) for x in first]
                elif "node_path" in obj and obj["node_path"] is not None:
                    single = obj["node_path"]
                    if isinstance(single, (list, tuple)):
                        node_path = [int(x) for x in single]

                self._records.append(SpanRecord(text=text, node_path=node_path, raw=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,
        )
        # drop batch dimension
        return {k: v.squeeze(0) for k, v in enc.items()}

    def _make_mlm_labels(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Very small, self-contained dynamic masking implementation.

        - 15% of non-special tokens are selected for prediction.
        - Of those, 80% replaced with [MASK], 10% with random token,
          10% left unchanged (BERT-style).
        """
        labels = input_ids.clone()

        if not hasattr(self.tokenizer, "get_special_tokens_mask"):
            # Fallback: treat everything as non-special
            special_tokens_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        else:
            special_tokens_mask = torch.tensor(
                self.tokenizer.get_special_tokens_mask(
                    input_ids.tolist(), already_has_special_tokens=True
                ),
                dtype=torch.bool,
            )

        # mask candidates: non-special tokens
        probability_matrix = torch.full(labels.shape, self.mask_probability)
        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)

        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100  # ignore index for unmasked tokens

        # 80% of selected tokens -> [MASK]
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        mask_token_id = self.tokenizer.mask_token_id
        if mask_token_id is not None:
            input_ids[indices_replaced] = mask_token_id

        # 10% of selected tokens -> random token
        indices_random = (
            torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
            & masked_indices
            & ~indices_replaced
        )
        random_words = torch.randint(
            low=0,
            high=self.tokenizer.vocab_size,
            size=labels.shape,
            dtype=torch.long,
        )
        input_ids[indices_random] = random_words[indices_random]

        # 10% remain unchanged
        return labels

    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))
        labels_mlm = self._make_mlm_labels(input_ids.clone())

        # Build per-level targets from canonical node_path if available.
        level_targets: Dict[int, int] = {}
        if rec.node_path is not None:
            for level, node_id in enumerate(rec.node_path):
                # Level 0 is typically the root; skip if you don't train on it.
                if level == 0:
                    continue
                level_targets[level] = int(node_id)

        sample: Dict[str, Any] = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels_mlm": labels_mlm,
            "level_targets": level_targets,
        }

        # Keep the full path information around for path-consistency and
        # contrastive losses. Prefer a richer `node_paths` field (multiple
        # valid paths for DAGs / multi-parent ontologies) if present; fall
        # back to a single `node_path` otherwise.
        paths = rec.raw.get("node_paths")
        if paths is None and rec.node_path is not None:
            paths = [list(rec.node_path)]
        if paths is not None:
            sample["paths"] = paths

        return sample

def collate_tree_of_life_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Custom collate_fn to handle:
      - dict-of-tensors for `level_targets`
      - list-of-paths for contrastive loss
    """
    # Simple stack for standard fields
    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_mlm = torch.stack([b["labels_mlm"] for b in batch], dim=0)

    # Merge level_targets into dict[level] -> tensor(batch,)
    # Collect all levels that appear in any sample.
    all_levels = set()
    for b in batch:
        all_levels.update(b["level_targets"].keys())

    level_targets: Dict[int, torch.Tensor] = {}
    for level in sorted(all_levels):
        targets_for_level = []
        for b in batch:
            # Use -100 for "unknown" / missing level for this sample
            targets_for_level.append(b["level_targets"].get(level, -100))
        level_targets[level] = torch.tensor(targets_for_level, dtype=torch.long)

    out: Dict[str, Any] = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels_mlm": labels_mlm,
        "level_targets": level_targets,
    }

    # Optional paths for contrastive loss
    if all("paths" in b for b in batch):
        out["paths"] = [b["paths"] for b in batch]

    return out


SpanRecord

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Any

import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizerBase

@dataclass
class SpanRecord:
    text: str
    # Canonical single path used for per-level classification targets.
    node_path: Optional[List[int]]
    # Full raw record, which may include richer DAG-style fields such as
    # `node_paths` (multiple valid paths) in addition to `node_path`.
    raw: Dict[str, Any]

class TreeOfLifeDataset(Dataset):
    """
    Minimal dataset for TOLBERT training.

    Assumes a JSONL `spans_file` where each line looks like:
        {
          "span_id": "s0",
          "text": "def forward(...):",
          "node_path": [0, 1, 2, 10]   # optional but recommended
        }

    You are free to extend this to include per-level fields instead of
    `node_path` (e.g., `level_1_id`, `level_2_id`, ...). This class is
    intentionally simple and meant as a starting point.
    """

    def __init__(
        self,


TreeOfLifeDataset

import json
from dataclasses import dataclass
from typing import Dict, List, Optional, Any

import torch
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizerBase

@dataclass
class SpanRecord:
    text: str
    # Canonical single path used for per-level classification targets.
    node_path: Optional[List[int]]
    # Full raw record, which may include richer DAG-style fields such as
    # `node_paths` (multiple valid paths) in addition to `node_path`.
    raw: Dict[str, Any]

class TreeOfLifeDataset(Dataset):
    """
    Minimal dataset for TOLBERT training.

    Assumes a JSONL `spans_file` where each line looks like:
        {
          "span_id": "s0",
          "text": "def forward(...):",
          "node_path": [0, 1, 2, 10]   # optional but recommended
        }

    You are free to extend this to include per-level fields instead of
    `node_path` (e.g., `level_1_id`, `level_2_id`, ...). This class is
    intentionally simple and meant as a starting point.
    """

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

        self._records: List[SpanRecord] = []
        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)
                text = obj["text"]

                # Support both legacy single-path and richer DAG-style formats:
                #   - "node_path": [v0, v1, ..., vK]
                #   - "node_paths": [[v0, v1, ..., vK], [v0, v1', ..., vK'], ...]
                #
                # We keep a single canonical path (first path if multiple are
                # provided) for per-level classification, while preserving the
                # richer structure in `raw` for path-consistency and
                # contrastive losses.
                node_path: Optional[List[int]] = None
                if "node_paths" in obj and obj["node_paths"] is not None:
                    paths_val = obj["node_paths"]
                    if isinstance(paths_val, list) and paths_val:
                        first = paths_val[0]
                        if isinstance(first, (list, tuple)):
                            node_path = [int(x) for x in first]
                elif "node_path" in obj and obj["node_path"] is not None:
                    single = obj["node_path"]
                    if isinstance(single, (list, tuple)):
                        node_path = [int(x) for x in single]

                self._records.append(SpanRecord(text=text, node_path=node_path, raw=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,
        )
        # drop batch dimension
        return {k: v.squeeze(0) for k, v in enc.items()}

    def _make_mlm_labels(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Very small, self-contained dynamic masking implementation.

        - 15% of non-special tokens are selected for prediction.
        - Of those, 80% replaced with [MASK], 10% with random token,
          10% left unchanged (BERT-style).
        """
        labels = input_ids.clone()

        if not hasattr(self.tokenizer, "get_special_tokens_mask"):
            # Fallback: treat everything as non-special
            special_tokens_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        else:
            special_tokens_mask = torch.tensor(
                self.tokenizer.get_special_tokens_mask(
                    input_ids.tolist(), already_has_special_tokens=True
                ),
                dtype=torch.bool,
            )

        # mask candidates: non-special tokens
        probability_matrix = torch.full(labels.shape, self.mask_probability)
        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)

        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100  # ignore index for unmasked tokens

        # 80% of selected tokens -> [MASK]
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        mask_token_id = self.tokenizer.mask_token_id
        if mask_token_id is not None:
            input_ids[indices_replaced] = mask_token_id

        # 10% of selected tokens -> random token
        indices_random = (
            torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
            & masked_indices
            & ~indices_replaced
        )
        random_words = torch.randint(
            low=0,
            high=self.tokenizer.vocab_size,
            size=labels.shape,
            dtype=torch.long,
        )
        input_ids[indices_random] = random_words[indices_random]

        # 10% remain unchanged
        return labels

    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))
        labels_mlm = self._make_mlm_labels(input_ids.clone())

        # Build per-level targets from canonical node_path if available.
        level_targets: Dict[int, int] = {}
        if rec.node_path is not None:
            for level, node_id in enumerate(rec.node_path):
                # Level 0 is typically the root; skip if you don't train on it.
                if level == 0:
                    continue
                level_targets[level] = int(node_id)

        sample: Dict[str, Any] = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels_mlm": labels_mlm,
            "level_targets": level_targets,
        }

        # Keep the full path information around for path-consistency and
        # contrastive losses. Prefer a richer `node_paths` field (multiple
        # valid paths for DAGs / multi-parent ontologies) if present; fall
        # back to a single `node_path` otherwise.
        paths = rec.raw.get("node_paths")
        if paths is None and rec.node_path is not None:
            paths = [list(rec.node_path)]
        if paths is not None:
            sample["paths"] = paths

        return sample

def collate_tree_of_life_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Custom collate_fn to handle:
      - dict-of-tensors for `level_targets`
      - list-of-paths for contrastive loss
    """
    # Simple stack for standard fields
    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_mlm = torch.stack([b["labels_mlm"] for b in batch], dim=0)

    # Merge level_targets into dict[level] -> tensor(batch,)
    # Collect all levels that appear in any sample.
    all_levels = set()
    for b in batch:
        all_levels.update(b["level_targets"].keys())

    level_targets: Dict[int, torch.Tensor] = {}


        sample: Dict[str, Any] = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels_mlm": labels_mlm,
            "level_targets": level_targets,
        }

        # Keep the full path information around for path-consistency and
        # contrastive losses. Prefer a richer `node_paths` field (multiple
        # valid paths for DAGs / multi-parent ontologies) if present; fall
        # back to a single `node_path` otherwise.
        paths = rec.raw.get("node_paths")
        if paths is None and rec.node_path is not None:
            paths = [list(rec.node_path)]
        if paths is not None:
            sample["paths"] = paths

        return sample

def collate_tree_of_life_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Custom collate_fn to handle:
      - dict-of-tensors for `level_targets`
      - list-of-paths for contrastive loss
    """
    # Simple stack for standard fields
    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_mlm = torch.stack([b["labels_mlm"] for b in batch], dim=0)

    # Merge level_targets into dict[level] -> tensor(batch,)
    # Collect all levels that appear in any sample.
    all_levels = set()
    for b in batch:
        all_levels.update(b["level_targets"].keys())

    level_targets: Dict[int, torch.Tensor] = {}
    for level in sorted(all_levels):
        targets_for_level = []
        for b in batch:
            # Use -100 for "unknown" / missing level for this sample
            targets_for_level.append(b["level_targets"].get(level, -100))
        level_targets[level] = torch.tensor(targets_for_level, dtype=torch.long)

    out: Dict[str, Any] = {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels_mlm": labels_mlm,
        "level_targets": level_targets,
    }

    # Optional paths for contrastive loss
    if all("paths" in b for b in batch):
        out["paths"] = [b["paths"] for b in batch]

    return out


    # `node_paths` (multiple valid paths) in addition to `node_path`.
    raw: Dict[str, Any]

class TreeOfLifeDataset(Dataset):
    """
    Minimal dataset for TOLBERT training.

    Assumes a JSONL `spans_file` where each line looks like:
        {
          "span_id": "s0",
          "text": "def forward(...):",
          "node_path": [0, 1, 2, 10]   # optional but recommended
        }

    You are free to extend this to include per-level fields instead of
    `node_path` (e.g., `level_1_id`, `level_2_id`, ...). This class is
    intentionally simple and meant as a starting point.
    """

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

        self._records: List[SpanRecord] = []
        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)
                text = obj["text"]

                # Support both legacy single-path and richer DAG-style formats:
                #   - "node_path": [v0, v1, ..., vK]
                #   - "node_paths": [[v0, v1, ..., vK], [v0, v1', ..., vK'], ...]
                #
                # We keep a single canonical path (first path if multiple are
                # provided) for per-level classification, while preserving the
                # richer structure in `raw` for path-consistency and
                # contrastive losses.
                node_path: Optional[List[int]] = None
                if "node_paths" in obj and obj["node_paths"] is not None:
                    paths_val = obj["node_paths"]
                    if isinstance(paths_val, list) and paths_val:
                        first = paths_val[0]
                        if isinstance(first, (list, tuple)):
                            node_path = [int(x) for x in first]
                elif "node_path" in obj and obj["node_path"] is not None:
                    single = obj["node_path"]
                    if isinstance(single, (list, tuple)):
                        node_path = [int(x) for x in single]

                self._records.append(SpanRecord(text=text, node_path=node_path, raw=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,
        )
        # drop batch dimension
        return {k: v.squeeze(0) for k, v in enc.items()}

    def _make_mlm_labels(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Very small, self-contained dynamic masking implementation.

        - 15% of non-special tokens are selected for prediction.


                #   - "node_paths": [[v0, v1, ..., vK], [v0, v1', ..., vK'], ...]
                #
                # We keep a single canonical path (first path if multiple are
                # provided) for per-level classification, while preserving the
                # richer structure in `raw` for path-consistency and
                # contrastive losses.
                node_path: Optional[List[int]] = None
                if "node_paths" in obj and obj["node_paths"] is not None:
                    paths_val = obj["node_paths"]
                    if isinstance(paths_val, list) and paths_val:
                        first = paths_val[0]
                        if isinstance(first, (list, tuple)):
                            node_path = [int(x) for x in first]
                elif "node_path" in obj and obj["node_path"] is not None:
                    single = obj["node_path"]
                    if isinstance(single, (list, tuple)):
                        node_path = [int(x) for x in single]

                self._records.append(SpanRecord(text=text, node_path=node_path, raw=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,
        )
        # drop batch dimension
        return {k: v.squeeze(0) for k, v in enc.items()}

    def _make_mlm_labels(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Very small, self-contained dynamic masking implementation.

        - 15% of non-special tokens are selected for prediction.
        - Of those, 80% replaced with [MASK], 10% with random token,
          10% left unchanged (BERT-style).
        """


                # provided) for per-level classification, while preserving the
                # richer structure in `raw` for path-consistency and
                # contrastive losses.
                node_path: Optional[List[int]] = None
                if "node_paths" in obj and obj["node_paths"] is not None:
                    paths_val = obj["node_paths"]
                    if isinstance(paths_val, list) and paths_val:
                        first = paths_val[0]
                        if isinstance(first, (list, tuple)):
                            node_path = [int(x) for x in first]
                elif "node_path" in obj and obj["node_path"] is not None:
                    single = obj["node_path"]
                    if isinstance(single, (list, tuple)):
                        node_path = [int(x) for x in single]

                self._records.append(SpanRecord(text=text, node_path=node_path, raw=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,
        )
        # drop batch dimension
        return {k: v.squeeze(0) for k, v in enc.items()}

    def _make_mlm_labels(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Very small, self-contained dynamic masking implementation.

        - 15% of non-special tokens are selected for prediction.
        - Of those, 80% replaced with [MASK], 10% with random token,
          10% left unchanged (BERT-style).
        """
        labels = input_ids.clone()

        if not hasattr(self.tokenizer, "get_special_tokens_mask"):
            # Fallback: treat everything as non-special
            special_tokens_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        else:
            special_tokens_mask = torch.tensor(
                self.tokenizer.get_special_tokens_mask(
                    input_ids.tolist(), already_has_special_tokens=True
                ),
                dtype=torch.bool,


                    single = obj["node_path"]
                    if isinstance(single, (list, tuple)):
                        node_path = [int(x) for x in single]

                self._records.append(SpanRecord(text=text, node_path=node_path, raw=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,
        )
        # drop batch dimension
        return {k: v.squeeze(0) for k, v in enc.items()}

    def _make_mlm_labels(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Very small, self-contained dynamic masking implementation.

        - 15% of non-special tokens are selected for prediction.
        - Of those, 80% replaced with [MASK], 10% with random token,
          10% left unchanged (BERT-style).
        """
        labels = input_ids.clone()

        if not hasattr(self.tokenizer, "get_special_tokens_mask"):
            # Fallback: treat everything as non-special
            special_tokens_mask = torch.zeros_like(input_ids, dtype=torch.bool)
        else:
            special_tokens_mask = torch.tensor(
                self.tokenizer.get_special_tokens_mask(
                    input_ids.tolist(), already_has_special_tokens=True
                ),
                dtype=torch.bool,
            )

        # mask candidates: non-special tokens
        probability_matrix = torch.full(labels.shape, self.mask_probability)
        probability_matrix.masked_fill_(special_tokens_mask, value=0.0)

        masked_indices = torch.bernoulli(probability_matrix).bool()
        labels[~masked_indices] = -100  # ignore index for unmasked tokens

        # 80% of selected tokens -> [MASK]
        indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
        mask_token_id = self.tokenizer.mask_token_id
        if mask_token_id is not None:
            input_ids[indices_replaced] = mask_token_id

        # 10% of selected tokens -> random token
        indices_random = (
            torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
            & masked_indices
            & ~indices_replaced
        )
        random_words = torch.randint(
            low=0,
            high=self.tokenizer.vocab_size,
            size=labels.shape,
            dtype=torch.long,
        )
        input_ids[indices_random] = random_words[indices_random]

        # 10% remain unchanged
        return labels

    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))
        labels_mlm = self._make_mlm_labels(input_ids.clone())

        # Build per-level targets from canonical node_path if available.
        level_targets: Dict[int, int] = {}
        if rec.node_path is not None:
            for level, node_id in enumerate(rec.node_path):
                # Level 0 is typically the root; skip if you don't train on it.
                if level == 0:
                    continue
                level_targets[level] = int(node_id)

        sample: Dict[str, Any] = {
            "input_ids": input_ids,


        if mask_token_id is not None:
            input_ids[indices_replaced] = mask_token_id

        # 10% of selected tokens -> random token
        indices_random = (
            torch.bernoulli(torch.full(labels.shape, 0.5)).bool()
            & masked_indices
            & ~indices_replaced
        )
        random_words = torch.randint(
            low=0,
            high=self.tokenizer.vocab_size,
            size=labels.shape,
            dtype=torch.long,
        )
        input_ids[indices_random] = random_words[indices_random]

        # 10% remain unchanged
        return labels

    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))
        labels_mlm = self._make_mlm_labels(input_ids.clone())

        # Build per-level targets from canonical node_path if available.
        level_targets: Dict[int, int] = {}
        if rec.node_path is not None:
            for level, node_id in enumerate(rec.node_path):
                # Level 0 is typically the root; skip if you don't train on it.
                if level == 0:
                    continue
                level_targets[level] = int(node_id)

        sample: Dict[str, Any] = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels_mlm": labels_mlm,
            "level_targets": level_targets,
        }

        # Keep the full path information around for path-consistency and
        # contrastive losses. Prefer a richer `node_paths` field (multiple
        # valid paths for DAGs / multi-parent ontologies) if present; fall
        # back to a single `node_path` otherwise.
        paths = rec.raw.get("node_paths")
        if paths is None and rec.node_path is not None:
            paths = [list(rec.node_path)]
        if paths is not None:
            sample["paths"] = paths

        return sample

def collate_tree_of_life_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Custom collate_fn to handle:
      - dict-of-tensors for `level_targets`
      - list-of-paths for contrastive loss
    """
    # Simple stack for standard fields
    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_mlm = torch.stack([b["labels_mlm"] for b in batch], dim=0)

    # Merge level_targets into dict[level] -> tensor(batch,)
    # Collect all levels that appear in any sample.
    all_levels = set()
    for b in batch:
        all_levels.update(b["level_targets"].keys())

    level_targets: Dict[int, torch.Tensor] = {}


tolbert.decoding

from __future__ import annotations

from typing import Dict, List, Optional, Sequence

import torch

def greedy_hierarchical_decode(
    level_logits: Dict[str, torch.Tensor],
    parent_to_children: Dict[int, Dict[int, List[int]]],
    levels: Optional[Sequence[int]] = None,
) -> Dict[int, torch.Tensor]:
    """
    Greedy top-down decoding of a hierarchical path with child masking.

    This helper mirrors the inference-time procedure described in the TOLBERT
    paper: for each level, predictions are restricted to children of the
    previously chosen parent whenever a parent->children mapping is provided.

    Args:
        level_logits:
            Mapping from level identifier (as used in TOLBERT.level_heads,
            typically the string form of an integer like \"1\", \"2\", ...) to
            a tensor of shape (batch, C_level) with logits for that level.
        parent_to_children:
            Mapping:
                level_int -> { parent_idx -> [child_idx, ...] }
            where indices are in the same class-index space as the logits for
            that level. For example, for level 2, parent indices live at
            level 1 and child indices live at level 2.
        levels:
            Optional explicit decoding order, e.g. [1, 2, 3]. If omitted, the
            function sorts the integer level keys found in `level_logits`.

    Returns:
        dict[level_int] -> tensor(batch,) of predicted class indices per level.

    Notes:
        - If no children are registered for a given (level, parent_idx), the
          decoder falls back to an unconstrained argmax over that level's
          logits for that example.
        - This helper does not assume anything about how node ids are
          assigned; it only assumes that the indices in `parent_to_children`
          are valid indices into the corresponding logits.
    """
    if not level_logits:
        return {}

    # Determine decoding order
    if levels is not None and len(levels) > 0:
        order = list(levels)
    else:
        order = sorted(int(k) for k in level_logits.keys())

    # Sanity: ensure we have logits for all requested levels
    order = [lvl for lvl in order if str(lvl) in level_logits]
    if not order:
        return {}

    batch_size = next(iter(level_logits.values())).size(0)
    device = next(iter(level_logits.values())).device

    preds: Dict[int, torch.Tensor] = {}

    # First level: plain argmax
    first_level = order[0]
    first_logits = level_logits[str(first_level)]
    preds[first_level] = first_logits.argmax(dim=-1)

    # Subsequent levels: mask by children whenever possible
    for idx in range(1, len(order)):
        level = order[idx]
        prev_level = order[idx - 1]

        logits = level_logits[str(level)]  # (B, C_level)
        parent_map = parent_to_children.get(level, {})

        # Default: unconstrained argmax for all examples
        level_preds = logits.argmax(dim=-1)

        if parent_map:
            # Refine predictions where we know the parent->children mapping.
            refined: List[int] = []
            parent_preds = preds[prev_level]
            for b in range(batch_size):
                parent_idx = int(parent_preds[b].item())
                children = parent_map.get(parent_idx)
                if not children:
                    # No constraints for this parent; keep unconstrained argmax.
                    refined.append(int(level_preds[b].item()))
                    continue

                # Restrict logits to the valid children for this parent.
                child_indices = torch.tensor(children, dtype=torch.long, device=device)
                # Guard against out-of-range indices.
                child_indices = child_indices[
                    (child_indices >= 0) & (child_indices < logits.size(-1))
                ]
                if child_indices.numel() == 0:
                    refined.append(int(level_preds[b].item()))
                    continue

                child_logits = logits[b, child_indices]
                rel_idx = int(torch.argmax(child_logits).item())
                refined.append(int(child_indices[rel_idx].item()))

            level_preds = torch.tensor(refined, device=device, dtype=torch.long)

        preds[level] = level_preds

    return preds


from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Tuple

# Basic identifiers
EntityId = str

@dataclass(frozen=True)
class Span:
    """Inclusive 1-based line span within an artifact."""

    start_line: int
    end_line: int

@dataclass
class Artifact:
    """
    Minimal file-level artifact in a program graph.

    - `uri` is typically a `program://` URI.
    - `type` is a coarse tag such as "source", "header", "binary".
    - `hash` is an implementation-defined content hash (may be empty).
    - `span` is optional; whole-file artifacts can leave it as None.
    """

    uri: str
    type: str
    hash: str
    span: Optional[Span] = None

@dataclass
class Entity:
    """
    Logical entity in a program graph.

    Common patterns:
      - Repository / project node.
      - File / module node.
      - Function / method / class node.

    `artifact_uri` and `span` are optional and only populated when the
    entity corresponds to a concrete region in a file.
    """

    id: EntityId
    kind: str
    uri: str
    artifact_uri: Optional[str] = None
    span: Optional[Span] = None
    labels: List[str] = field(default_factory=list)
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class Edge:
    """
    Typed relationship between entities.

    Typical edge types:
      - "owns": containment (repo→file, file→function, module→class, ...)
      - "imports": dependency between modules / files / packages
      - "calls": function/method call
      - "tests": test module or case covering a target
      - "similar_to": embedding- or heuristic-based similarity
    """

    src: EntityId
    dst: EntityId
    type: str
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class ResolvedAnchor:
    """
    Resolution of a URI back onto a concrete artifact region.
    """

    artifact_uri: str
    span: Span
    hash: str

class ProgramGraph:
    """
    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


Span

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Tuple

# Basic identifiers
EntityId = str

@dataclass(frozen=True)
class Span:
    """Inclusive 1-based line span within an artifact."""

    start_line: int
    end_line: int

@dataclass
class Artifact:
    """
    Minimal file-level artifact in a program graph.

    - `uri` is typically a `program://` URI.
    - `type` is a coarse tag such as "source", "header", "binary".
    - `hash` is an implementation-defined content hash (may be empty).
    - `span` is optional; whole-file artifacts can leave it as None.
    """

    uri: str
    type: str
    hash: str
    span: Optional[Span] = None

@dataclass


Artifact

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Tuple

# Basic identifiers
EntityId = str

@dataclass(frozen=True)
class Span:
    """Inclusive 1-based line span within an artifact."""

    start_line: int
    end_line: int

@dataclass
class Artifact:
    """
    Minimal file-level artifact in a program graph.

    - `uri` is typically a `program://` URI.
    - `type` is a coarse tag such as "source", "header", "binary".
    - `hash` is an implementation-defined content hash (may be empty).
    - `span` is optional; whole-file artifacts can leave it as None.
    """

    uri: str
    type: str
    hash: str
    span: Optional[Span] = None

@dataclass
class Entity:
    """
    Logical entity in a program graph.

    Common patterns:
      - Repository / project node.
      - File / module node.
      - Function / method / class node.

    `artifact_uri` and `span` are optional and only populated when the
    entity corresponds to a concrete region in a file.
    """

    id: EntityId
    kind: str
    uri: str
    artifact_uri: Optional[str] = None


Entity

@dataclass
class Artifact:
    """
    Minimal file-level artifact in a program graph.

    - `uri` is typically a `program://` URI.
    - `type` is a coarse tag such as "source", "header", "binary".
    - `hash` is an implementation-defined content hash (may be empty).
    - `span` is optional; whole-file artifacts can leave it as None.
    """

    uri: str
    type: str
    hash: str
    span: Optional[Span] = None

@dataclass
class Entity:
    """
    Logical entity in a program graph.

    Common patterns:
      - Repository / project node.
      - File / module node.
      - Function / method / class node.

    `artifact_uri` and `span` are optional and only populated when the
    entity corresponds to a concrete region in a file.
    """

    id: EntityId
    kind: str
    uri: str
    artifact_uri: Optional[str] = None
    span: Optional[Span] = None
    labels: List[str] = field(default_factory=list)
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class Edge:
    """
    Typed relationship between entities.

    Typical edge types:
      - "owns": containment (repo→file, file→function, module→class, ...)
      - "imports": dependency between modules / files / packages
      - "calls": function/method call
      - "tests": test module or case covering a target
      - "similar_to": embedding- or heuristic-based similarity
    """

    src: EntityId
    dst: EntityId
    type: str
    attributes: Dict[str, Any] = field(default_factory=dict)


Edge

    Common patterns:
      - Repository / project node.
      - File / module node.
      - Function / method / class node.

    `artifact_uri` and `span` are optional and only populated when the
    entity corresponds to a concrete region in a file.
    """

    id: EntityId
    kind: str
    uri: str
    artifact_uri: Optional[str] = None
    span: Optional[Span] = None
    labels: List[str] = field(default_factory=list)
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class Edge:
    """
    Typed relationship between entities.

    Typical edge types:
      - "owns": containment (repo→file, file→function, module→class, ...)
      - "imports": dependency between modules / files / packages
      - "calls": function/method call
      - "tests": test module or case covering a target
      - "similar_to": embedding- or heuristic-based similarity
    """

    src: EntityId
    dst: EntityId
    type: str
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class ResolvedAnchor:
    """
    Resolution of a URI back onto a concrete artifact region.
    """

    artifact_uri: str
    span: Span
    hash: str

class ProgramGraph:
    """
    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships


ResolvedAnchor

@dataclass
class Edge:
    """
    Typed relationship between entities.

    Typical edge types:
      - "owns": containment (repo→file, file→function, module→class, ...)
      - "imports": dependency between modules / files / packages
      - "calls": function/method call
      - "tests": test module or case covering a target
      - "similar_to": embedding- or heuristic-based similarity
    """

    src: EntityId
    dst: EntityId
    type: str
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class ResolvedAnchor:
    """
    Resolution of a URI back onto a concrete artifact region.
    """

    artifact_uri: str
    span: Span
    hash: str

class ProgramGraph:
    """
    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []


ProgramGraph

      - "similar_to": embedding- or heuristic-based similarity
    """

    src: EntityId
    dst: EntityId
    type: str
    attributes: Dict[str, Any] = field(default_factory=dict)

@dataclass
class ResolvedAnchor:
    """
    Resolution of a URI back onto a concrete artifact region.
    """

    artifact_uri: str
    span: Span
    hash: str

class ProgramGraph:
    """
    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


entities

    span: Span
    hash: str

class ProgramGraph:
    """
    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


edges

class ProgramGraph:
    """
    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


artifacts

    Abstract, language-agnostic view over a program / repository.

    Implementations are expected to provide at least:
      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


resolve

      - entities(): iterable of Entity nodes
      - edges(): iterable of Edge relationships
      - artifacts(): file-level view
      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


      - resolve(): URI → concrete anchor

    Subclasses may add richer APIs as needed.
    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self


subgraph

    """

    # The default implementation is intentionally skeletal and returns
    # empty views; concrete graphs should override these.

    def entities(self) -> Iterable[Entity]:
        return []

    def edges(self) -> Iterable[Edge]:
        return []

    def artifacts(self, kind: str) -> Iterable[Artifact]:
        return []

    def resolve(self, uri: str) -> ResolvedAnchor:
        raise NotImplementedError("resolve() must be implemented by subclasses")

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

    def subgraph(self, seeds: List[EntityId], radius: int) -> "ProgramGraph":
        return self