| """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: |
| 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__}") |
|
|
