presentation/src/common.py

"""Shared helpers for the BITS Phishing Detector pipeline.

Owned by the data-prep stage; every other src/*.py imports it as ``import common``.
Only stdlib + yaml + numpy + torch are touched here so the module is cheap to import
from CLI entry points, notebooks and tests alike.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import platform
import random
import re
import sys
import timeit
import unicodedata
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np
import yaml

_PATH_KEYS = frozenset(
    {"data_raw", "data_processed", "tokenizer", "checkpoints", "reports", "logs", "app_dir"}
)
_NON_ALNUM = re.compile(r"[^a-z0-9]+")


def project_root() -> Path:
    """Parent of ``src/``; every relative path in config.yaml is anchored here."""
    return Path(__file__).resolve().parent.parent


def load_config(path: str | Path = "config.yaml") -> dict[str, Any]:
    """yaml.safe_load of the project config with ``paths.*`` resolved against the project root.

    Resolving at load time means every script can be launched from any cwd and still write
    to the same data/, checkpoints/ and logs/ directories.
    """
    cfg_path = Path(path)
    if not cfg_path.is_absolute():
        cfg_path = project_root() / cfg_path
    with cfg_path.open("r", encoding="utf-8") as fh:
        cfg = yaml.safe_load(fh)
    paths = cfg.get("paths") or {}
    for key, value in list(paths.items()):
        if key in _PATH_KEYS and isinstance(value, str) and not os.path.isabs(value):
            paths[key] = str(project_root() / value)
    cfg["paths"] = paths
    cfg["_config_path"] = str(cfg_path)
    return cfg


def set_seed(seed: int) -> None:
    """Seed random, numpy and torch (CPU + every CUDA device).

    Deliberately does NOT call ``torch.use_deterministic_algorithms``: that would disable the
    fast SDPA / cuBLAS kernels the 337M encoder depends on, and bitwise reproducibility is not a
    project requirement (seeded splits + logged configs are).
    """
    random.seed(seed)
    np.random.seed(seed)
    try:
        import torch

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:  # pragma: no cover - torch is pinned, but keep common.py importable without it
        pass


def setup_logging(name: str, log_dir: str | Path = "logs") -> logging.Logger:
    """Console + file logger (``logs/<name>-<YYYYmmdd-HHMMSS>.log``) at INFO.

    Each script gets its own timestamped file so that re-runs never clobber the evidence of an
    earlier run; RUN_LOG.md quotes from these files.
    """
    log_path = Path(log_dir)
    if not log_path.is_absolute():
        log_path = project_root() / log_path
    log_path.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    logfile = log_path / f"{name}-{stamp}.log"

    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)
    logger.propagate = False
    for handler in list(logger.handlers):
        logger.removeHandler(handler)
    fmt = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s", "%H:%M:%S")
    stream = logging.StreamHandler(sys.stdout)
    stream.setFormatter(fmt)
    file_handler = logging.FileHandler(logfile, encoding="utf-8")
    file_handler.setFormatter(fmt)
    logger.addHandler(stream)
    logger.addHandler(file_handler)
    logger.info("log file: %s", logfile)
    return logger


def print_device_info() -> dict[str, Any]:
    """Print and return the runtime facts that matter for VRAM planning; CPU-safe."""
    info: dict[str, Any] = {
        "python": platform.python_version(),
        "torch": None,
        "cuda_build": None,
        "cudnn": None,
        "device": "cpu",
        "sm": None,
        "vram_gib": 0.0,
        "bf16": False,
        "free_gib": 0.0,
    }
    try:
        import torch

        info["torch"] = torch.__version__
        info["cuda_build"] = torch.version.cuda
        info["cudnn"] = torch.backends.cudnn.version()
        if torch.cuda.is_available():
            idx = torch.cuda.current_device()
            props = torch.cuda.get_device_properties(idx)
            free, total = torch.cuda.mem_get_info(idx)
            info.update(
                device=props.name,
                sm=f"{props.major}.{props.minor}",
                vram_gib=round(total / 2**30, 2),
                bf16=bool(torch.cuda.is_bf16_supported()),
                free_gib=round(free / 2**30, 2),
            )
    except ImportError:
        pass
    for key, value in info.items():
        print(f"{key:>10}: {value}")
    return info


def vram_stats() -> dict[str, float]:
    """Current/peak allocated and reserved VRAM in GB (zeros on CPU)."""
    zeros = {"allocated_gb": 0.0, "reserved_gb": 0.0, "max_allocated_gb": 0.0, "max_reserved_gb": 0.0}
    try:
        import torch

        if not torch.cuda.is_available():
            return zeros
        gb = float(2**30)
        return {
            "allocated_gb": torch.cuda.memory_allocated() / gb,
            "reserved_gb": torch.cuda.memory_reserved() / gb,
            "max_allocated_gb": torch.cuda.max_memory_allocated() / gb,
            "max_reserved_gb": torch.cuda.max_memory_reserved() / gb,
        }
    except ImportError:
        return zeros


def timer() -> float:
    return timeit.default_timer()


def normalize_text(s: str) -> str:
    """Dedup normalisation: NFKC -> lowercase -> keep only [a-z0-9].

    Removing every non-alphanumeric character neutralises the Enron-Spam " , " tokenisation,
    whitespace/line-ending differences and HTML entity noise, so two renderings of the same email
    collide on the same key. The NFKC step is an extension of the verify_data.md section 7 wording
    ("lowercase, keep [a-z0-9]"): it folds full-width and ligature variants onto the same key and
    only ever merges more rows, so counts differ marginally from that document (e.g. 16,260 vs
    16,259 TREC-06 rows under a body-only filter). Every sha1 key in data/processed derives from
    THIS function; downstream code (app/, notebooks) must import or copy it verbatim, never
    re-derive it from the one-line description.
    """
    if not s:
        return ""
    return _NON_ALNUM.sub("", unicodedata.normalize("NFKC", s).lower())


def sha1_hex(s: str) -> str:
    return hashlib.sha1(s.encode("utf-8", "surrogatepass")).hexdigest()


def write_json(obj: Any, path: str | Path) -> None:
    p = Path(path)
    p.parent.mkdir(parents=True, exist_ok=True)
    with p.open("w", encoding="utf-8") as fh:
        json.dump(obj, fh, indent=2, sort_keys=False, default=_json_default)
        fh.write("\n")


def read_json(path: str | Path) -> Any:
    with Path(path).open("r", encoding="utf-8") as fh:
        return json.load(fh)


def append_run_log(markdown: str, path: str | Path = "RUN_LOG.md") -> None:
    """Append a blank line + the markdown block to RUN_LOG.md (the project's running record)."""
    p = Path(path)
    if not p.is_absolute():
        p = project_root() / p
    with p.open("a", encoding="utf-8") as fh:
        fh.write("\n" + markdown.rstrip("\n") + "\n")


def _json_default(obj: Any) -> Any:
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, Path):
        return str(obj)
    if isinstance(obj, datetime):
        return obj.isoformat()
    raise TypeError(f"not JSON serialisable: {type(obj).__name__}")