File size: 3,015 Bytes
877049d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | 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(),
}
|