src/utils.py

import json
import logging
import os
import random
import re
from pathlib import Path
from typing import Any, Dict

import numpy as np
import torch
import yaml


LABEL_TO_ID = {
    "non-clickbait": 0,
    "clickbait": 1,
}

ID_TO_LABEL = {value: key for key, value in LABEL_TO_ID.items()}


def ensure_dir(path: str | Path) -> Path:
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    return path


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


def clean_title(text: Any) -> str:
    value = "" if text is None else str(text)
    value = re.sub(r"\s+", " ", value).strip()
    return value


def normalize_label(label: Any) -> int:
    if isinstance(label, (int, np.integer)):
        if int(label) in ID_TO_LABEL:
            return int(label)
    normalized = str(label).strip().lower()
    if normalized not in LABEL_TO_ID:
        raise ValueError(f"Unsupported label value: {label}")
    return LABEL_TO_ID[normalized]


def load_yaml(path: str | Path) -> Dict[str, Any]:
    with Path(path).open("r", encoding="utf-8") as file:
        return yaml.safe_load(file) or {}


def save_json(data: Dict[str, Any], path: str | Path) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as file:
        json.dump(data, file, indent=2, ensure_ascii=False)


def resolve_device(preferred: str | None = None) -> torch.device:
    if preferred:
        return torch.device(preferred)
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def compute_class_weights(labels: np.ndarray | list[int]) -> torch.Tensor:
    values = np.asarray(labels, dtype=np.int64)
    counts = np.bincount(values)
    counts[counts == 0] = 1
    weights = values.shape[0] / (len(counts) * counts)
    return torch.tensor(weights, dtype=torch.float32)


def setup_logger(output_dir: str | Path, name: str = "train") -> logging.Logger:
    output_dir = ensure_dir(output_dir)
    log_path = output_dir / f"{name}.log"
    logger = logging.getLogger(str(log_path))
    logger.setLevel(logging.INFO)
    logger.handlers.clear()
    formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(message)s")

    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)
    return logger


def count_parameters(model: torch.nn.Module) -> int:
    return sum(param.numel() for param in model.parameters() if param.requires_grad)


def environment_summary() -> Dict[str, Any]:
    return {
        "pythonhashseed": os.environ.get("PYTHONHASHSEED"),
        "cuda_available": torch.cuda.is_available(),
        "device_count": torch.cuda.device_count(),
    }