src/lineage.py

"""Model lineage, walked over the `base_model` column.

Every model in the index may declare one base model. That makes the index a
forest, and this module is how the Space reads it: ancestors of a model,
descendants of a model, and the biggest family in the index (the design's
"lineage spotlight", which on real Hub data is the FinBERT tree).

**Every walk here is cycle-safe.** `base_model` is author-declared metadata,
not a verified relation, and the Hub does not stop anyone from declaring a
cycle -- A says it is finetuned from B while B says it is finetuned from A.
Two models that each claim the other would spin an unguarded walk forever and
hang the Space. So every traversal carries a `seen` set and stops the moment it
revisits a node, and every walk is additionally depth-capped.

The design draws at most a handful of generations, so MAX_DEPTH is a display
bound as much as a safety one.
"""

from __future__ import annotations

MAX_DEPTH = 24


def build_parents(rows) -> dict:
    """{model_id: declared base model} for rows whose parent is in the index.

    A base model that is not itself indexed is kept -- the design shows
    "fine-tuned from X" even when X is a general-purpose model outside the
    Atlas -- but it terminates the walk, since there is no row to step to.
    """
    return {
        row["id"]: (row.get("base_model") or "")
        for row in rows
        if row.get("id")
    }


def build_children(parents: dict) -> dict:
    """Invert the parent map: {model_id: [direct children]}."""
    children: dict = {}
    for child, parent in parents.items():
        if not parent or parent == child:
            continue
        children.setdefault(parent, []).append(child)
    for kids in children.values():
        kids.sort()
    return children


def ancestors(model_id: str, parents: dict, max_depth: int = MAX_DEPTH) -> list:
    """Walk up the tree from `model_id`, nearest parent first.

    Stops on: no declared parent, a parent outside the index, the depth cap,
    or a node already visited. That last condition is what makes a circular
    declaration terminate instead of hanging.
    """
    chain = []
    seen = {model_id}
    current = model_id
    for _ in range(max_depth):
        parent = parents.get(current) or ""
        if not parent or parent in seen:
            break
        chain.append(parent)
        seen.add(parent)
        if parent not in parents:
            # Declared but not indexed: it is a real ancestor and worth
            # showing, but there is nothing further to walk to.
            break
        current = parent
    return chain


def root_of(model_id: str, parents: dict, max_depth: int = MAX_DEPTH) -> str:
    chain = ancestors(model_id, parents, max_depth)
    return chain[-1] if chain else model_id


def descendants(model_id: str, children: dict, max_depth: int = MAX_DEPTH) -> list:
    """Every model below `model_id`, breadth-first, each returned once.

    Breadth-first so the design's tree lists direct children before
    grandchildren. `seen` guards both diamonds (two paths to one node) and
    cycles.
    """
    out = []
    seen = {model_id}
    frontier = [model_id]
    for _ in range(max_depth):
        following = []
        for node in frontier:
            for child in children.get(node, ()):
                if child in seen:
                    continue
                seen.add(child)
                out.append(child)
                following.append(child)
        if not following:
            break
        frontier = following
    return out


def direct_children(model_id: str, children: dict) -> list:
    return list(children.get(model_id, ()))


def siblings(model_id: str, parents: dict, children: dict) -> list:
    """Other models sharing this model's declared parent."""
    parent = parents.get(model_id) or ""
    if not parent:
        return []
    return [c for c in children.get(parent, ()) if c != model_id]


def largest_family(rows, parents: dict = None, children: dict = None):
    """The indexed model with the most descendants: (root_id, [descendants]).

    This is what fills the design's "lineage spotlight" panel. Ties break on
    model id so the panel does not reshuffle between identical rebuilds.
    """
    rows = list(rows)
    if parents is None:
        parents = build_parents(rows)
    if children is None:
        children = build_children(parents)

    indexed = {row["id"] for row in rows if row.get("id")}
    best_id, best_kids = "", []
    for candidate in sorted(indexed):
        kids = descendants(candidate, children)
        if len(kids) > len(best_kids):
            best_id, best_kids = candidate, kids
    return best_id, best_kids