utils/modular_model_detector.py

# Copyright 2025 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# 🔴🔴🔴 THIS IS AN INTERNAL TOOL. It WILL interact with the hub and use significant local compute resources. Use at your own risk.


"""
Modular model detector: utilities for detecting code similarities between model implementations.

This module provides tools to analyze and detect similarities between different model implementations
in the transformers library. It uses both embedding-based and token-based (Jaccard) similarity metrics
to identify similar code patterns across different model definitions.

Its function is to identify which models can be _modular_-ized, meaning, which already existing classes are
present in the codebase and look very similar to the one we have.

Two scores are computed, one is a code embedding, and the other is a simple Jaccard bag-of-tokens index for overlap
of token sets. A score of 1.00 means the code is identical.

Usage:

```bash
cd transformers

# Use directly the util, it will download the index embedding from the hub. It will require some RAM/VRAM.

>>> python utils/modular_model_detector.py --modeling-file my_new_beit3_modeling_file.py
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 33.62it/s]
encoding 21 query definitions with Qwen/Qwen3-Embedding-4B (device=cuda, batch=16, max_length=4096)
stuff.py::Beit3ImageTextMatchingOutput:
embedding:
    blip_2::Blip2ImageTextMatchingModelOutput (0.9994)
    chinese_clip::ChineseCLIPOutput (0.9818)
    owlvit::OwlViTOutput (0.9818)
    aimv2::Aimv2Output (0.9818)
    blip::BlipOutput (0.9818)
jaccard:
    owlv2::Owlv2Output (0.9667)
    metaclip_2::MetaClip2Output (0.9667)
    altclip::AltCLIPOutput (0.9667)
    owlvit::OwlViTOutput (0.9667)
    blip::BlipOutput (0.9667)
intersection:
    blip::BlipOutput
    owlvit::OwlViTOutput

stuff.py::Beit3MLP:
embedding:
    efficientloftr::EfficientLoFTRMLP (0.9718)
    seggpt::SegGptMlp (0.9650)
    mgp_str::MgpstrMlp (0.9646)
    vitpose_backbone::VitPoseBackboneMLP (0.9640)
    granitemoeshared::GraniteMoeSharedMLP (0.9633)
jaccard:
    chinese_clip::ChineseCLIPTextSelfOutput (0.5294)
    convbert::ConvBertSelfOutput (0.5294)
    bert::BertSelfOutput (0.5294)
    roformer::RoFormerSelfOutput (0.5294)
    layoutlmv3::LayoutLMv3SelfOutput (0.5294)
intersection:

stuff.py::Beit3FeedForwardNetwork:
embedding:
    prophetnet::ProphetNetFeedForward (0.9766)
    dab_detr::DabDetrDecoderLayerFFN (0.9730)
    kosmos2::Kosmos2TextFFN (0.9697)
    kosmos2_5::Kosmos2_5TextFFN (0.9697)
    parakeet::ParakeetEncoderFeedForward (0.9678)
jaccard:
    groupvit::GroupViTMLP (0.4898)
    convbert::ConvBertOutput (0.4600)
    chinese_clip::ChineseCLIPTextOutput (0.4565)
    bert::BertOutput (0.4565)
    roformer::RoFormerOutput (0.4565)
intersection:



```


# If you wish to build the index first, you can run

python utils/modular_model_detector.py --build

# You can also change the embedding model for a larger/smaller one.
"""

import argparse
import ast
import json
import logging
import os
import re
from datetime import datetime
from functools import cache
from pathlib import Path

import numpy as np
import torch
from huggingface_hub import HfApi, snapshot_download
from huggingface_hub import logging as huggingface_hub_logging
from safetensors.numpy import load_file as safetensors_load
from safetensors.numpy import save_file as safetensors_save
from tqdm import tqdm

import transformers
from transformers import AutoModel, AutoTokenizer
from transformers.utils import enable_tf32
from transformers.utils import logging as transformers_logging


# ANSI color codes for CLI output styling
ANSI_RESET = "\033[0m"
ANSI_BOLD = "\033[1m"
ANSI_HEADER = "\033[1;36m"
ANSI_SECTION = "\033[1;35m"
ANSI_ROW = "\033[0;37m"
ANSI_HIGHLIGHT_TOP = "\033[1;32m"
ANSI_HIGHLIGHT_OLD = "\033[1;33m"
ANSI_HIGHLIGHT_CANDIDATE = "\033[1;34m"


os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["TRANSFORMERS_VERBOSITY"] = "error"

MODELS_ROOT = Path("src/transformers/models")
EMBEDDINGS_PATH = "embeddings.safetensors"
INDEX_MAP_PATH = "code_index_map.json"
TOKENS_PATH = "code_index_tokens.json"
HUB_DATASET_DEFAULT = "hf-internal-testing/transformers_code_embeddings"

EMBEDDING_MODEL = "Qwen/Qwen3-Embedding-4B"
BATCH_SIZE = 16
MAX_LENGTH = 4096


def _normalize(string: str | None) -> str:
    """
    Normalize a string by removing all non-alphanumeric characters and converting to lowercase.

    Args:
        string (`str` or `None`): The string to normalize.

    Returns:
        `str`: The normalized string, or empty string if input is None.
    """
    return re.sub(r"[^a-z0-9]+", "", string.lower()) if string else ""


def _strip_source_for_tokens(code: str) -> str:
    """
    Strip docstrings, comments, and import statements from source code.

    Args:
        code (`str`): The source code to strip.

    Returns:
        `str`: The stripped source code.
    """
    code = re.sub(r'("""|\'\'\')(?:.|\n)*?\1', "", code)
    code = re.sub(r"#.*", "", code)
    return "\n".join(line for line in code.splitlines() if not re.match(r"\s*(from|import)\s+", line))


def _tokenize(code: str) -> set[str]:
    """
    Extract all Python identifiers from source code.

    Args:
        code (`str`): The source code to tokenize.

    Returns:
        `set[str]`: A set of all identifiers found in the code.
    """
    return set(re.findall(r"\b[a-zA-Z_][a-zA-Z0-9_]*\b", code))


def _leading_symbol_prefix(name: str) -> str:
    """
    Extract the leading prefix from a symbol name (e.g., 'Llama' from 'LlamaAttention').

    Args:
        name (`str`): The symbol name to extract prefix from.

    Returns:
        `str`: The leading prefix, or empty string if no match.
    """
    match = re.match(r"^([A-Z][a-z0-9]+)", name) or re.match(r"^([A-Za-z0-9]+)", name)
    return match.group(1) if match else ""


def _sanitize_for_embedding(code: str, model_hint: str | None, symbol_hint: str | None) -> str:
    """
    Sanitize code for embedding by replacing model-specific identifiers with generic placeholder.

    Args:
        code (`str`): The source code to sanitize.
        model_hint (`str` or `None`): Hint about the model name (e.g., 'llama').
        symbol_hint (`str` or `None`): Hint about the symbol name (e.g., 'LlamaAttention').

    Returns:
        `str`: The sanitized code with model-specific identifiers replaced by 'Model'.
    """
    base = _strip_source_for_tokens(code)
    variants = set()
    if model_hint:
        variants.add(model_hint)
        variants.add(model_hint.replace("_", ""))
        variants.add(re.sub(r"\d+", "", model_hint))
    if symbol_hint:
        prefix = _leading_symbol_prefix(symbol_hint)
        if prefix:
            variants.add(prefix)
            variants.add(prefix.replace("_", ""))
            variants.add(re.sub(r"\d+", "", prefix))
    variants |= {variant.lower() for variant in list(variants)}
    sanitized = base
    for variant in sorted({x for x in variants if len(x) >= 3}, key=len, reverse=True):
        sanitized = re.sub(re.escape(variant), "Model", sanitized, flags=re.IGNORECASE)
    return sanitized


class CodeSimilarityAnalyzer:
    """
    Analyzer for detecting code similarities between model implementations.

    This class uses embedding-based and token-based similarity metrics to identify similar
    code patterns across different model definitions in the transformers library.

    Args:
        hub_dataset (`str`): The Hub dataset repository ID containing the code embeddings index.
    """

    def __init__(self, hub_dataset: str):
        for name in ("huggingface_hub", "httpx", "urllib3", "transformers"):
            logging.getLogger(name).setLevel(logging.ERROR)
        huggingface_hub_logging.set_verbosity_error()
        transformers_logging.set_verbosity_error()
        enable_tf32(True)
        torch.set_grad_enabled(False)

        self.models_root = MODELS_ROOT
        self.hub_dataset = hub_dataset
        self.tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL)
        self.model = AutoModel.from_pretrained(EMBEDDING_MODEL, torch_dtype="auto", device_map="auto").eval()

        self.device = self.model.device
        self.index_dir: Path | None = None

    # ---------- HUB IO ----------

    def _resolve_index_path(self, filename: str) -> Path:
        if self.index_dir is None:
            return Path(filename)
        return self.index_dir / filename

    def ensure_local_index(self) -> None:
        """Ensure index files are available locally, preferring Hub cache snapshots."""
        if self.index_dir is not None and all(
            (self.index_dir / fname).exists() for fname in (EMBEDDINGS_PATH, INDEX_MAP_PATH, TOKENS_PATH)
        ):
            return

        workspace_dir = Path.cwd()
        if all((workspace_dir / fname).exists() for fname in (EMBEDDINGS_PATH, INDEX_MAP_PATH, TOKENS_PATH)):
            self.index_dir = workspace_dir
            return

        logging.info(f"downloading index from hub cache: {self.hub_dataset}")
        snapshot_path = snapshot_download(repo_id=self.hub_dataset, repo_type="dataset")
        snapshot_dir = Path(snapshot_path)
        missing = [
            fname for fname in (EMBEDDINGS_PATH, INDEX_MAP_PATH, TOKENS_PATH) if not (snapshot_dir / fname).exists()
        ]
        if missing:
            raise FileNotFoundError("Missing expected files in Hub snapshot: " + ", ".join(missing))
        self.index_dir = snapshot_dir

    def push_index_to_hub(self) -> None:
        """Upload index files to the Hub dataset repository."""
        api = HfApi()
        api.create_repo(repo_id=self.hub_dataset, repo_type="dataset", exist_ok=True)
        for fname in (EMBEDDINGS_PATH, INDEX_MAP_PATH, TOKENS_PATH):
            logging.info(f"pushing {fname} -> {self.hub_dataset}")
            api.upload_file(
                path_or_fileobj=fname,
                path_in_repo=os.path.basename(fname),
                repo_id=self.hub_dataset,
                repo_type="dataset",
            )

    # ---------- parsing & encoding ----------

    def _extract_definitions(
        self, file_path: Path, relative_to: Path | None = None, model_hint: str | None = None
    ) -> tuple[dict[str, str], dict[str, str], dict[str, list[str]], dict[str, str]]:
        """
        Extract class and function definitions from a Python file.

        Args:
            file_path (`Path`): Path to the Python file to parse.
            relative_to (`Path` or `None`): Base path for computing relative identifiers.
            model_hint (`str` or `None`): Model name hint for sanitization.

        Returns:
            `tuple[dict[str, str], dict[str, str], dict[str, list[str]], dict[str, str]]`: A tuple containing:
                - definitions_raw: Mapping of identifiers to raw source code
                - definitions_sanitized: Mapping of identifiers to sanitized source code
                - definitions_tokens: Mapping of identifiers to sorted token lists
                - definitions_kind: Mapping of identifiers to either "class" or "function"
        """
        definitions_raw = {}
        definitions_sanitized = {}
        definitions_tokens = {}
        definitions_kind = {}
        source = file_path.read_text(encoding="utf-8")
        lines = source.splitlines()
        tree = ast.parse(source)
        for node in ast.iter_child_nodes(tree):
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
                segment = ast.get_source_segment(source, node)
                if segment is None and hasattr(node, "lineno") and hasattr(node, "end_lineno"):
                    start = max(0, node.lineno - 1)
                    end = node.end_lineno
                    segment = "\n".join(lines[start:end])
                if segment:
                    identifier = (
                        f"{file_path.relative_to(relative_to)}:{node.name}"
                        if relative_to
                        else f"{file_path.name}:{node.name}"
                    )
                    definitions_raw[identifier] = segment
                    sanitized = _sanitize_for_embedding(segment, model_hint, node.name)
                    definitions_sanitized[identifier] = sanitized
                    definitions_tokens[identifier] = sorted(_tokenize(sanitized))
                    if isinstance(node, ast.ClassDef):
                        definitions_kind[identifier] = "class"
                    else:
                        definitions_kind[identifier] = "function"
        return definitions_raw, definitions_sanitized, definitions_tokens, definitions_kind

    def _infer_model_from_relative_path(self, relative_path: Path) -> str | None:
        try:
            relative = relative_path.resolve().relative_to(self.models_root.resolve())
            return relative.parts[0]
        except Exception:
            return None

    def _infer_query_model_name(self, modeling_file: Path) -> str | None:
        model = self._infer_model_from_relative_path(modeling_file)
        if model:
            return model
        stem = modeling_file.stem
        if stem.startswith("modeling_") and len(stem) > len("modeling_"):
            return stem[len("modeling_") :]
        return None

    def _encode_batch(self, texts: list[str]) -> np.ndarray:
        """
        Encode a batch of texts into normalized embeddings.

        Args:
            texts (`list[str]`): List of text strings to encode.

        Returns:
            `np.ndarray`: Normalized embeddings as a float32 numpy array.
        """
        encoded = self.tokenizer(texts, padding=True, truncation=True, max_length=MAX_LENGTH, return_tensors="pt")
        encoded = {key: value.to(self.device) for key, value in encoded.items()}
        with (
            torch.autocast(device_type=self.device.type, dtype=self.dtype)
            if self.device.type == "cuda"
            else torch.no_grad()
        ):
            output = self.model(**encoded)
            if hasattr(output, "last_hidden_state"):
                embeddings = output.last_hidden_state
                mask = encoded["attention_mask"].unsqueeze(-1)
                embeddings = (embeddings * mask).sum(dim=1) / mask.sum(dim=1).clamp_min(1e-9)
            elif hasattr(output, "pooler_output"):
                embeddings = output.pooler_output
            else:
                embeddings = output[0].mean(dim=1)
        embeddings = torch.nn.functional.normalize(embeddings.float(), p=2, dim=1)
        return embeddings.cpu().numpy().astype("float32")

    def encode(self, texts: list[str]) -> np.ndarray:
        """
        Encode a list of texts into embeddings, processing in batches.

        Args:
            texts (`list[str]`): List of text strings to encode.

        Returns:
            `np.ndarray`: Stacked embeddings for all texts.
        """
        output = []
        for i in tqdm(range(0, len(texts), BATCH_SIZE), desc="encode", leave=False):
            output.append(self._encode_batch(texts[i : i + BATCH_SIZE]))
            if self.device.type == "cuda":
                torch.cuda.empty_cache()
        return np.vstack(output) if output else np.zeros((0, 0), dtype="float32")

    # ---------- build & search ----------

    def build_index(self) -> None:
        """Build the code similarity index from all modeling files and save to disk."""
        logging.info("collecting files")
        files = list(self.models_root.rglob("modeling_*.py"))
        logging.info(f"parsing {len(files)} files")

        identifiers = []
        sanitized_sources = []
        tokens_map = {}

        for file_path in tqdm(files, desc="parse", leave=False):
            model_hint = self._infer_model_from_relative_path(file_path)
            (
                _,
                definitions_sanitized,
                definitions_tokens,
                _,
            ) = self._extract_definitions(file_path, self.models_root, model_hint)
            for identifier in definitions_sanitized.keys():
                identifiers.append(identifier)
                sanitized_sources.append(definitions_sanitized[identifier])
                tokens_map[identifier] = definitions_tokens[identifier]

        logging.info(
            f"encoding {len(sanitized_sources)} definitions with {EMBEDDING_MODEL} (device={self.device.type}, batch={BATCH_SIZE}, max_length={MAX_LENGTH})"
        )
        embeddings = self.encode(sanitized_sources)
        safetensors_save({"embeddings": embeddings}, EMBEDDINGS_PATH)
        with open(INDEX_MAP_PATH, "w", encoding="utf-8") as file:
            json.dump({int(i): identifiers[i] for i in range(len(identifiers))}, file)
        with open(TOKENS_PATH, "w", encoding="utf-8") as file:
            json.dump(tokens_map, file)

        self.index_dir = Path.cwd()

    def _topk_embedding(
        self,
        query_embedding_row: np.ndarray,
        base_embeddings: np.ndarray,
        identifier_map: dict[int, str],
        self_model_normalized: str,
        self_name: str,
        k: int,
    ) -> list[tuple[str, float]]:
        similarities = query_embedding_row @ base_embeddings.T
        indices = np.argpartition(-similarities, k + 32)[: k + 32]
        indices = indices[np.argsort(-similarities[indices])]
        output = []
        for match_id in indices:
            identifier = identifier_map[int(match_id)]
            parent_relative_path, match_name = identifier.split(":", 1)
            parent_model = Path(parent_relative_path).parts[0]
            if match_name == self_name:
                continue
            if self_model_normalized and _normalize(parent_model) == self_model_normalized:
                continue
            output.append((identifier, float(similarities[match_id])))
            if len(output) >= k:
                break
        return output

    def _topk_jaccard(
        self,
        query_tokens: set[str],
        identifiers: list[str],
        tokens_map: dict[str, list[str]],
        self_model_normalized: str,
        self_name: str,
        k: int,
    ) -> list[tuple[str, float]]:
        """
        Find top-k most similar definitions using Jaccard similarity on token sets.

        Args:
            query_tokens (`set[str]`): Set of tokens from the query definition.
            identifiers (`list[str]`): List of all definition identifiers in the index.
            tokens_map (`dict[str, list[str]]`): Mapping of identifiers to their token lists.
            self_model_normalized (`str`): Normalized name of the query model to exclude.
            self_name (`str`): Name of the query definition to exclude.
            k (`int`): Number of top results to return.

        Returns:
            `list[tuple[str, float]]`: List of (identifier, score) tuples.
        """
        scores = []
        for identifier in identifiers:
            parent_relative_path, match_name = identifier.split(":", 1)
            parent_model = Path(parent_relative_path).parts[0]
            if match_name == self_name:
                continue
            if self_model_normalized and _normalize(parent_model) == self_model_normalized:
                continue
            tokens = set(tokens_map.get(identifier, []))
            if not tokens or not query_tokens:
                continue
            score = len(query_tokens & tokens) / len(query_tokens | tokens)
            if score > 0:
                scores.append((identifier, score))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:k]

    def analyze_file(
        self, modeling_file: Path, top_k_per_item: int = 5, allow_hub_fallback: bool = True, use_jaccard=False
    ) -> dict[str, dict[str, list]]:
        """
        Analyze a modeling file and find similar code definitions in the index.

        Args:
            modeling_file (`Path`): Path to the modeling file to analyze.
            top_k_per_item (`int`, *optional*, defaults to 5): Number of top matches to return per definition.
            allow_hub_fallback (`bool`, *optional*, defaults to `True`): Whether to download index from Hub if not found locally.

        Returns:
            `dict[str, dict[str, list]]`: Dictionary mapping definition names to their similarity results.
                Each result contains 'embedding', 'jaccard', and 'intersection' keys.
        """
        if allow_hub_fallback:
            self.ensure_local_index()

        base = safetensors_load(str(self._resolve_index_path(EMBEDDINGS_PATH)))
        base_embeddings = base["embeddings"]
        with open(self._resolve_index_path(INDEX_MAP_PATH), "r", encoding="utf-8") as file:
            identifier_map = {int(key): value for key, value in json.load(file).items()}
        identifiers = [identifier_map[i] for i in range(len(identifier_map))]
        with open(self._resolve_index_path(TOKENS_PATH), "r", encoding="utf-8") as file:
            tokens_map = json.load(file)

        self_model = self._infer_query_model_name(modeling_file)
        definitions_raw, definitions_sanitized, _, definitions_kind = self._extract_definitions(
            modeling_file, None, self_model
        )
        query_identifiers = list(definitions_raw.keys())
        query_sources_sanitized = [definitions_sanitized[key] for key in query_identifiers]
        query_tokens_list = [set(_tokenize(source)) for source in query_sources_sanitized]
        self_model_normalized = _normalize(self_model)

        logging.info(
            f"encoding {len(query_sources_sanitized)} query definitions with {EMBEDDING_MODEL} (device={self.device.type}, batch={BATCH_SIZE}, max_length={MAX_LENGTH})"
        )
        query_embeddings = self.encode(query_sources_sanitized)

        output = {}
        for i, query_identifier in enumerate(query_identifiers):
            query_name = query_identifier.split(":")[-1]
            embedding_top = self._topk_embedding(
                query_embeddings[i], base_embeddings, identifier_map, self_model_normalized, query_name, top_k_per_item
            )
            embedding_set = {identifier for identifier, _ in embedding_top}
            kind = definitions_kind.get(query_identifier, "function")
            entry = {"kind": kind, "embedding": embedding_top}
            if use_jaccard:
                jaccard_top = self._topk_jaccard(
                    query_tokens_list[i], identifiers, tokens_map, self_model_normalized, query_name, top_k_per_item
                )
                jaccard_set = {identifier for identifier, _ in jaccard_top}
                intersection = set(embedding_set & jaccard_set)

                entry.update({"jaccard": jaccard_top, "intersection": intersection})
            output[query_name] = entry
        return output


_RELEASE_RE = re.compile(
    r"(?:^|[\*_`\s>])(?:this|the)\s+model\s+was\s+released\s+on\s+(\d{4}-\d{2}-\d{2})\b", re.IGNORECASE
)


def build_date_data() -> dict[str, str]:
    """
    Scan Markdown files in `root_dir` and build {model_id: date_released}.

    - model_id is the filename without extension (e.g., "llama" for "llama.md")
    - date_released is the first YYYY-MM-DD matched after "...was released on ..."
    - Ignores non-*.md files and directories.

    Returns:
        dict[str, str]: mapping of model_id -> ISO date string (YYYY-MM-DD).
                        Files without a match are simply omitted.
    """

    root_dir = transformers.__file__.split("src/transformers")[0]
    root = Path(root_dir).joinpath("docs/source/en/model_doc")
    result: dict[str, str] = {}

    for md_path in root.glob("*.md"):
        try:
            text = md_path.read_text(encoding="utf-8", errors="ignore")
        except Exception:
            # Skip unreadable files quietly
            logging.info(f"Failed to read md for {md_path}")

        m = _RELEASE_RE.search(text)
        if m:
            model_id = md_path.stem  # e.g., "llama" from "llama.md"
            result[model_id] = m.group(1)

    return result


def _format_table(headers: list[str], rows: list[tuple[str, ...] | None], row_styles: list[str] | None = None) -> str:
    if not rows:
        return f"{ANSI_ROW}(no matches){ANSI_RESET}"

    widths = [len(header) for header in headers]
    for row in rows:
        if row is None:
            continue
        for idx, cell in enumerate(row):
            widths[idx] = max(widths[idx], len(cell))

    header_line = " | ".join(header.ljust(widths[idx]) for idx, header in enumerate(headers))
    divider = "-+-".join("-" * widths[idx] for idx in range(len(headers)))
    total_width = sum(widths) + 3 * (len(headers) - 1)

    styled_rows = []
    style_idx = 0
    for row in rows:
        if row is None:
            styled_rows.append(f"{ANSI_SECTION}{'-' * total_width}{ANSI_RESET}")
            continue

        line = " | ".join(cell.ljust(widths[col_idx]) for col_idx, cell in enumerate(row))
        style = ANSI_ROW
        if row_styles and style_idx < len(row_styles) and row_styles[style_idx]:
            style = row_styles[style_idx]
        styled_rows.append(f"{style}{line}{ANSI_RESET}")
        style_idx += 1

    return "\n".join([f"{ANSI_SECTION}{header_line}{ANSI_RESET}", divider] + styled_rows)


def _parse_release_date(value: str) -> datetime | None:
    """Return a datetime parsed from YYYY-MM-DD strings, otherwise None."""
    try:
        return datetime.strptime(value, "%Y-%m-%d")
    except (TypeError, ValueError):
        return None


@cache
def _load_definition_line_map(relative_path: str) -> dict[str, int]:
    """Return {definition_name: line_number} for top-level definitions in the given file."""
    file_path = MODELS_ROOT / relative_path
    try:
        source = file_path.read_text(encoding="utf-8")
    except (FileNotFoundError, OSError):
        return {}  # gracefully keep going

    try:
        tree = ast.parse(source)
    except SyntaxError:
        return {}

    line_map: dict[str, int] = {}
    for node in ast.iter_child_nodes(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            line_map[node.name] = getattr(node, "lineno", None) or 1
        elif isinstance(node, ast.Assign):
            continue
    return line_map


def _resolve_definition_location(relative_path: str, definition: str) -> tuple[str, str]:
    """Return full path and formatted line number string for the given definition."""
    full_path = MODELS_ROOT / relative_path
    line = _load_definition_line_map(relative_path).get(definition)
    line_str = str(line) if line is not None else "?"
    return str(full_path), line_str


def _colorize_heading(text: str) -> str:
    return f"{ANSI_HEADER}{ANSI_BOLD}{text}{ANSI_RESET}"


def main():
    """CLI entry point for the modular model detector."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    parser = argparse.ArgumentParser(prog="hf-code-sim")
    parser.add_argument("--build", action="store_true")
    parser.add_argument("--modeling-file", type=str, help='You can just specify "vits" if you are lazy like me.')
    parser.add_argument(
        "--push-new-index", action="store_true", help="After --build, push index files to a Hub dataset."
    )
    parser.add_argument(
        "--hub-dataset", type=str, default=HUB_DATASET_DEFAULT, help="Hub dataset repo id to pull/push the index."
    )
    parser.add_argument("--use_jaccard", type=bool, default=False, help="Whether or not to use jaccard index")
    args = parser.parse_args()

    analyzer = CodeSimilarityAnalyzer(hub_dataset=args.hub_dataset)

    if args.build:
        analyzer.build_index()
        if args.push_new_index:
            analyzer.push_index_to_hub()
        return

    if not args.modeling_file:
        raise SystemExit("Provide --modeling-file or use --build")

    dates = build_date_data()
    modeling_file = args.modeling_file
    if os.sep not in modeling_file:
        modeling_file = os.path.join("src", "transformers", "models", modeling_file, f"modeling_{modeling_file}.py")

    results = analyzer.analyze_file(
        Path(modeling_file), top_k_per_item=5, allow_hub_fallback=True, use_jaccard=args.use_jaccard
    )
    modeling_filename = Path(modeling_file).name
    release_key = modeling_filename.split("modeling_")[-1][:-3]
    release_date = dates.get(release_key, "unknown release date")

    aggregate_scores: dict[str, float] = {}
    for data in results.values():
        for identifier, score in data.get("embedding", []):
            try:
                relative_path, _ = identifier.split(":", 1)
            except ValueError:
                continue
            aggregate_scores[relative_path] = aggregate_scores.get(relative_path, 0.0) + score

    best_candidate_path: str | None = None
    if aggregate_scores:
        best_candidate_path = max(aggregate_scores.items(), key=lambda item: item[1])[0]
        best_model = Path(best_candidate_path).parts[0] if Path(best_candidate_path).parts else "?"
        best_release = dates.get(best_model, "unknown release date")
        logging.info(
            f"{ANSI_HIGHLIGHT_CANDIDATE}Closest overall candidate: {MODELS_ROOT / best_candidate_path}"
            f" (release: {best_release}, total score: {aggregate_scores[best_candidate_path]:.4f}){ANSI_RESET}"
        )

    grouped: dict[str, list[tuple[str, dict]]] = {"class": [], "function": []}
    for query_name, data in results.items():
        kind = data.get("kind", "function")
        grouped.setdefault(kind, []).append((query_name, data))

    section_titles = [("class", "Classes"), ("function", "Functions")]
    legend_shown = False
    for kind, title in section_titles:
        entries = grouped.get(kind, [])
        if not entries:
            continue

        metrics_present: set[str] = set()
        for _, data in entries:
            if data.get("embedding"):
                metrics_present.add("embedding")
            if args.use_jaccard:
                if data.get("jaccard"):
                    metrics_present.add("jaccard")
                if data.get("intersection"):
                    metrics_present.add("intersection")

        include_metric_column = bool(metrics_present - {"embedding"})
        headers = ["Symbol", "Path", "Score", "Release"]
        if include_metric_column:
            headers = ["Symbol", "Metric", "Path", "Score", "Release"]

        table_rows: list[tuple[str, ...] | None] = []
        row_styles: list[str] = []
        has_metric_rows = False

        logging.info(_colorize_heading(title))

        for query_name, data in entries:
            if table_rows:
                table_rows.append(None)

            symbol_label = query_name
            if release_date:
                symbol_label = f"{symbol_label}"

            symbol_row = (symbol_label,) + ("",) * (len(headers) - 1)
            table_rows.append(symbol_row)
            row_styles.append(ANSI_BOLD)

            embedding_details: list[tuple[str, str, str, float, str]] = []
            embedding_style_indices: list[int] = []

            for identifier, score in data.get("embedding", []):
                try:
                    relative_path, match_name = identifier.split(":", 1)
                except ValueError:
                    continue
                model_id = Path(relative_path).parts[0] if Path(relative_path).parts else "?"
                match_release = dates.get(model_id, "unknown release date")
                full_path, line = _resolve_definition_location(relative_path, match_name)
                display_path = f"{full_path}:{line} ({match_name})"

                if include_metric_column:
                    row = ("", "embedding", display_path, f"{score:.4f}", match_release)
                else:
                    row = ("", display_path, f"{score:.4f}", match_release)

                table_rows.append(row)
                row_styles.append(ANSI_ROW)
                embedding_style_indices.append(len(row_styles) - 1)
                embedding_details.append((relative_path, model_id, match_name, score, match_release))
                has_metric_rows = True

            if embedding_details:
                highest_score = None
                highest_idx = None
                for idx, (_, _, _, score, _) in enumerate(embedding_details):
                    if highest_score is None or score > highest_score:
                        highest_score = score
                        highest_idx = idx

                if highest_idx is not None:
                    row_styles[embedding_style_indices[highest_idx]] = ANSI_HIGHLIGHT_TOP

                if highest_score is not None:
                    oldest_idx = None
                    oldest_date = None
                    for idx, (_, model_id, _, score, release_value) in enumerate(embedding_details):
                        if highest_score - score > 0.1:
                            continue
                        parsed = _parse_release_date(release_value)
                        if parsed is None:
                            continue
                        if oldest_date is None or parsed < oldest_date:
                            oldest_date = parsed
                            oldest_idx = idx
                    if (
                        oldest_idx is not None
                        and row_styles[embedding_style_indices[oldest_idx]] != ANSI_HIGHLIGHT_TOP
                    ):
                        row_styles[embedding_style_indices[oldest_idx]] = ANSI_HIGHLIGHT_OLD

                if best_candidate_path is not None:
                    for idx, (relative_path, _, _, _, _) in enumerate(embedding_details):
                        style_position = embedding_style_indices[idx]
                        if row_styles[style_position] != ANSI_ROW:
                            continue
                        if relative_path == best_candidate_path:
                            row_styles[style_position] = ANSI_HIGHLIGHT_CANDIDATE

            if args.use_jaccard:
                for identifier, score in data.get("jaccard", []):
                    try:
                        relative_path, match_name = identifier.split(":", 1)
                    except ValueError:
                        continue
                    model_id = Path(relative_path).parts[0] if Path(relative_path).parts else "?"
                    match_release = dates.get(model_id, "unknown release date")
                    full_path, line = _resolve_definition_location(relative_path, match_name)
                    display_path = f"{full_path}:{line} ({match_name})"

                    if include_metric_column:
                        row = ("", "jaccard", display_path, f"{score:.4f}", match_release)
                    else:
                        row = ("", display_path, f"{score:.4f}", match_release)

                    table_rows.append(row)
                    row_styles.append(ANSI_ROW)
                    has_metric_rows = True
                    if best_candidate_path == relative_path:
                        row_styles[-1] = ANSI_HIGHLIGHT_CANDIDATE

                for identifier in sorted(data.get("intersection", [])):
                    try:
                        relative_path, match_name = identifier.split(":", 1)
                    except ValueError:
                        continue
                    model_id = Path(relative_path).parts[0] if Path(relative_path).parts else "?"
                    match_release = dates.get(model_id, "unknown release date")
                    full_path, line = _resolve_definition_location(relative_path, match_name)
                    display_path = f"{full_path}:{line} ({match_name})"

                    if include_metric_column:
                        row = ("", "intersection", display_path, "--", match_release)
                    else:
                        row = ("", display_path, "--", match_release)

                    table_rows.append(row)
                    row_styles.append(ANSI_ROW)
                    has_metric_rows = True
                    if best_candidate_path == relative_path:
                        row_styles[-1] = ANSI_HIGHLIGHT_CANDIDATE

        if table_rows:
            if not legend_shown and has_metric_rows:
                logging.info(
                    "Legend: "
                    f"{ANSI_HIGHLIGHT_TOP}highest match{ANSI_RESET}, "
                    f"{ANSI_HIGHLIGHT_OLD}oldest within 0.1{ANSI_RESET}, "
                    f"{ANSI_HIGHLIGHT_CANDIDATE}closest overall candidate{ANSI_RESET}"
                )
                legend_shown = True

            logging.info(_format_table(headers, table_rows, row_styles))
            logging.info("")


if __name__ == "__main__":
    main()