model/pico_type/labels.py

"""pico-type: label vocabularies and output decoder."""

from __future__ import annotations

from typing import Any, Dict, Optional

import torch


COARSE_LABELS = [
    "text", "code", "link", "image", "file", "config",
    "markup", "data", "error", "secret", "archive", "binary",
]

MODALITY_LABELS = [
    "textual", "binary_image", "binary_archive", "binary_executable",
    "binary_document", "binary_audio", "binary_video", "binary_other",
]

SUBTYPE_LABELS = [
    "json", "yaml", "toml", "ini", "csv", "tsv", "xml",
    "html", "markdown", "rst", "asciidoc", "tex",
    "sql", "graphql", "protobuf", "msgpack",
    "log", "diff", "patch", "env",
    "shell", "makefile", "dockerfile", "gitignore",
]

CODE_LANG_LABELS = [
    "python", "javascript", "typescript", "jsx", "tsx",
    "java", "kotlin", "scala", "groovy", "clojure",
    "c", "cpp", "csharp", "fsharp", "objectivec",
    "go", "rust", "zig",
    "ruby", "php", "perl", "lua", "tcl",
    "swift", "dart", "julia", "nim", "crystal",
    "haskell", "ocaml", "elm", "erlang", "elixir",
    "lisp", "scheme", "racket",
    "r", "matlab", "octave", "sas", "stata",
    "sql", "plsql", "tsql",
    "html", "css", "scss", "sass", "less",
    "bash", "zsh", "fish", "powershell",
    "vim", "fortran", "cobol", "ada", "pascal",
    "delphi", "vb", "prolog", "vhdl",
]

assert len(CODE_LANG_LABELS) == 62, f"expected 62 code langs, got {len(CODE_LANG_LABELS)}"

TEXT_LANG_LABELS = [
    "en", "es", "fr", "de", "it", "pt", "nl", "sv", "no", "da",
    "fi", "pl", "cs", "sk", "hu", "ro", "el", "tr",
    "ru", "uk", "bg", "sr", "hr",
    "zh", "ja", "ko", "vi", "th", "id", "ms",
]

assert len(TEXT_LANG_LABELS) == 30, f"expected 30 text langs, got {len(TEXT_LANG_LABELS)}"

FILE_MIME_LABELS = [
    "application/pdf", "application/zip", "application/gzip", "application/x-tar",
    "application/x-7z-compressed", "application/x-rar-compressed", "application/x-bzip2",
    "application/x-xz", "application/json", "application/xml", "application/yaml",
    "application/octet-stream", "application/x-executable", "application/x-mach-binary",
    "application/x-elf", "application/x-deb", "application/x-rpm",
    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "application/vnd.openxmlformats-officedocument.presentationml.presentation",
    "application/vnd.ms-excel", "application/vnd.ms-powerpoint",
    "application/msword", "application/rtf", "application/epub+zip",
    "application/x-ndjson",
    "text/plain", "text/csv", "text/html", "text/xml", "text/markdown",
    "image/png", "image/jpeg", "image/gif", "image/webp", "image/svg+xml",
    "image/bmp", "image/tiff", "image/heic", "image/heif", "image/avif",
    "image/x-icon", "image/vnd.adobe.photoshop",
    "video/mp4", "video/webm", "video/x-matroska", "video/quicktime",
    "video/x-msvideo", "video/x-flv", "video/x-mpeg",
    "audio/mpeg", "audio/ogg", "audio/wav", "audio/flac", "audio/aac",
    "audio/x-m4a", "audio/webm", "audio/midi",
    "font/ttf", "font/otf", "font/woff", "font/woff2",
    "application/x-sqlite3", "application/x-parquet",
    "application/x-protobuf", "application/x-flatbuffers",
    "application/x-cpio", "application/x-iso9660-image",
    "application/vnd.android.package-archive", "application/x-jar",
    "application/x-python-bytecode", "application/x-archive",
    "application/pgp-encrypted", "application/pgp-signature",
    "application/x-x509-ca-cert", "application/x-pem-file",
    "application/vnd.tcpdump.pcap",
    "application/java-vm",
    "application/x-matlab-data",
    "application/x-shockwave-flash",
    "application/x-font-ttf", "application/x-font-otf",
    "application/wasm", "application/x-ruby",
    "application/javascript", "application/ecmascript",
    "application/x-bittorrent", "application/x-dvi",
    "chemical/x-mdl-sdfile",
    "application/x-lzma",
]

assert len(FILE_MIME_LABELS) == 90, f"expected 90 file mimes, got {len(FILE_MIME_LABELS)}"

RISK_LABELS = ["api_key", "jwt", "ssh_key", "password", "email", "phone"]

assert len(RISK_LABELS) == 6, f"expected 6 risk labels, got {len(RISK_LABELS)}"

UNDETECTED = "undetected"

ALL_HEADS = ("coarse", "modality", "subtype", "code_lang", "text_lang", "file_mime", "risk")

HEAD_NUM_CLASSES: Dict[str, int] = {
    "coarse": len(COARSE_LABELS),
    "modality": len(MODALITY_LABELS),
    "subtype": len(SUBTYPE_LABELS),
    "code_lang": len(CODE_LANG_LABELS),
    "text_lang": len(TEXT_LANG_LABELS),
    "file_mime": len(FILE_MIME_LABELS),
    "risk": len(RISK_LABELS),
}

SUBTYPE_GATED_BY = {"config", "markup", "data"}
CODE_LANG_GATED_BY = {"code"}
TEXT_LANG_GATED_BY = {"text"}
FILE_MIME_GATED_BY = {
    "image", "file", "archive", "binary",
    "binary_image", "binary_archive", "binary_executable",
    "binary_document", "binary_audio", "binary_video", "binary_other",
}


def _pick_with_threshold(
    logits: torch.Tensor,
    labels: list[str],
    threshold: float,
) -> tuple[Optional[str], float]:
    probs = torch.softmax(logits, dim=-1)
    conf, idx = probs.max(dim=-1)
    if conf.item() < threshold:
        return UNDETECTED, conf.item()
    return labels[idx.item()], conf.item()


def decode_output(
    logits: Dict[str, torch.Tensor],
    tier: str = "base",
    undetected_threshold: float = 0.4,
    risk_threshold: float = 0.5,
) -> Dict[str, Any]:
    coarse_idx = int(logits["coarse"].argmax(dim=-1).item())
    coarse = COARSE_LABELS[coarse_idx]
    coarse_conf = float(torch.softmax(logits["coarse"], dim=-1).max().item())

    modality_idx = int(logits["modality"].argmax(dim=-1).item())
    modality = MODALITY_LABELS[modality_idx]
    modality_conf = float(torch.softmax(logits["modality"], dim=-1).max().item())

    subtype: Optional[str] = None
    if coarse in SUBTYPE_GATED_BY:
        subtype, _ = _pick_with_threshold(
            logits["subtype"], SUBTYPE_LABELS, undetected_threshold,
        )

    code_lang: Optional[str] = None
    if coarse in CODE_LANG_GATED_BY:
        code_lang, _ = _pick_with_threshold(
            logits["code_lang"], CODE_LANG_LABELS, undetected_threshold,
        )

    text_lang: Optional[str] = None
    if coarse in TEXT_LANG_GATED_BY:
        text_lang, _ = _pick_with_threshold(
            logits["text_lang"], TEXT_LANG_LABELS, undetected_threshold,
        )

    file_mime: Optional[str] = None
    if coarse in FILE_MIME_GATED_BY or modality in FILE_MIME_GATED_BY:
        file_mime, _ = _pick_with_threshold(
            logits["file_mime"], FILE_MIME_LABELS, undetected_threshold,
        )

    risk_logits = logits["risk"]
    risk_probs = torch.sigmoid(risk_logits).squeeze(0)
    risk_flags = [
        RISK_LABELS[i] for i in range(len(RISK_LABELS))
        if float(risk_probs[i].item()) >= risk_threshold
    ]

    return {
        "coarse": coarse,
        "modality": modality,
        "subtype": subtype,
        "code_language": code_lang,
        "text_language": text_lang,
        "file_mime": file_mime,
        "risk_flags": risk_flags,
        "confidence": coarse_conf,
        "modality_confidence": modality_conf,
        "model_tier": tier,
    }


def label_for(head: str, idx: int) -> str:
    table = {
        "coarse": COARSE_LABELS,
        "modality": MODALITY_LABELS,
        "subtype": SUBTYPE_LABELS,
        "code_lang": CODE_LANG_LABELS,
        "text_lang": TEXT_LANG_LABELS,
        "file_mime": FILE_MIME_LABELS,
        "risk": RISK_LABELS,
    }
    return table[head][idx]