BioinstLab/gmass-demo

1r"""2core.utils — Shared I/O, caching, and environment helpers.3MediSafe-GH · Biomedical Technologies Lab4 5Unified from two parallel implementations (Team D scratch work + the6GMASS_Coding_Standard.md reference repo). Function names from BOTH7versions are kept as aliases so nothing else in the codebase breaks:8 9    load_jsonl()         — returns [] on missing file (does not raise)10    append_jsonl()        \__ same function, two names11    save_jsonl_line()     /12    load_completed_ids()  — works with either function name above13 14Key cost-saving utility: load_completed_ids() enables crash-safe resumption15of API batches — never re-pay for a probe already evaluated.16"""17 18import json19import os20import re21import time22from datetime import datetime, timezone23from pathlib import Path24 25from core.logger import get_logger26 27logger = get_logger(__name__)28 29 30# ══════════════════════════════════════════════════════════════════════════════31# JSONL I/O32# ══════════════════════════════════════════════════════════════════════════════33 34def load_jsonl(path: str | Path, warn_missing: bool = True) -> list[dict]:35    """36    Load all records from a JSONL file.37 38    Returns [] if the file is missing (does not raise) — this matches the39    coding-standard reference implementation. Batch scripts can call this40    on an output file that doesn't exist yet without wrapping in try/except.41    """42    p = Path(path)43    if not p.exists():44        message = f"JSONL not found: {p} - returning []"45        if warn_missing:46            logger.warning(message)47        else:48            logger.debug(message)49        return []50 51    records, errors = [], 052    with open(p, encoding="utf-8") as f:53        for i, line in enumerate(f, 1):54            line = line.strip()55            if not line:56                continue57            try:58                records.append(json.loads(line))59            except json.JSONDecodeError as e:60                logger.error(f"JSON error line {i} of {p}: {e}")61                errors += 162 63    if errors:64        logger.warning(f"Loaded {len(records)} records with {errors} parse errors from {p}")65    else:66        logger.debug(f"Loaded {len(records)} records from {p}")67    return records68 69 70def append_jsonl(record: dict, path: str | Path) -> None:71    """72    Append ONE record to a JSONL file (creates file + parent dirs if missing).73 74    Always append — never overwrite — during batch runs. A crashed run loses75    at most one in-flight record, not the entire batch.76    """77    p = Path(path)78    p.parent.mkdir(parents=True, exist_ok=True)79    with open(p, "a", encoding="utf-8") as f:80        f.write(json.dumps(record, ensure_ascii=False) + "\n")81 82 83# Alias — earlier pipeline.py / scorer.py scripts call this name.84# Keeping both names means neither version of the codebase needs editing.85save_jsonl_line = append_jsonl86 87 88def load_completed_ids(output_path: str | Path, id_field: str = "probe_id") -> set[str]:89    """90    Return the set of probe_ids already present in an output JSONL.91 92    Use this at the start of every batch run to skip already-evaluated probes.93    This is the primary API cost-saving mechanism: zero re-calls on resume.94 95    Example:96        done   = load_completed_ids("data/eval_outputs/raw/gpt-4o.jsonl")97        probes = [p for p in all_probes if p["probe_id"] not in done]98        logger.info(f"Resuming: {len(probes)} probes remaining")99    """100    records = load_jsonl(output_path, warn_missing=False)101    ids = {r[id_field] for r in records if id_field in r}102    if ids:103        logger.info(f"Resume: {len(ids)} probes already done in {Path(output_path).name}")104    return ids105 106 107# ══════════════════════════════════════════════════════════════════════════════108# ENVIRONMENT DETECTION109# ══════════════════════════════════════════════════════════════════════════════110 111def is_kaggle() -> bool:112    """True when running inside a Kaggle kernel."""113    return os.getenv("KAGGLE_KERNEL_RUN_TYPE") is not None114 115 116def is_cuda_available() -> bool:117    """True when a CUDA GPU (RTX) is available."""118    try:119        import torch120        return torch.cuda.is_available()121    except ImportError:122        return False123 124 125def get_device() -> str:126    """Return 'cuda' on RTX, else 'cpu'. Used by local model loaders (LlamaGuard3, RoBERTa)."""127    return "cuda" if is_cuda_available() else "cpu"128 129 130def log_environment(logger_instance) -> None:131    """Log a one-line environment summary at run start."""132    if is_kaggle():133        env = "Kaggle (T4 GPU)" if is_cuda_available() else "Kaggle (CPU)"134    elif is_cuda_available():135        try:136            import torch137            name = torch.cuda.get_device_name(0)138            env  = f"Local CUDA — {name}"139        except Exception:140            env = "Local CUDA"141    else:142        env = "CPU only"143    logger_instance.info(f"Environment: {env}")144 145 146# ══════════════════════════════════════════════════════════════════════════════147# API HELPERS148# ══════════════════════════════════════════════════════════════════════════════149 150def get_api_key(env_var: str) -> str:151    """152    Retrieve API key from environment. Raises a clear ValueError if missing.153    Never hardcode keys in scripts — always use .env + dotenv.154    """155    key = os.getenv(env_var)156    if not key:157        raise ValueError(158            f"'{env_var}' not set. Add it to your .env file and call load_dotenv()."159        )160    return key161 162 163def retry_with_backoff(fn, retries: int = 3, base_wait: float = 1.0):164    """165    Call fn() with exponential backoff on RateLimitError.166    Returns None on final failure — never crashes the batch.167 168    Args:169        fn        : zero-argument callable (lambda wrapping the API call)170        retries   : max attempts171        base_wait : initial wait in seconds (doubles each retry)172    """173    for attempt in range(retries):174        try:175            return fn()176        except Exception as e:177            err_str = str(e).lower()178            is_rate = any(x in err_str for x in ["rate limit", "429", "quota"])179            wait = base_wait * (2 ** attempt)180            if is_rate:181                logger.warning(f"Rate limit hit (attempt {attempt+1}/{retries}). Waiting {wait}s.")182                time.sleep(wait)183            else:184                logger.error(f"API error (attempt {attempt+1}/{retries}): {e}")185                if attempt == retries - 1:186                    return None187                time.sleep(wait)188    return None189 190 191# ══════════════════════════════════════════════════════════════════════════════192# MISC193# ══════════════════════════════════════════════════════════════════════════════194 195def utc_now() -> str:196    """Return current UTC time as ISO-8601 string."""197    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")198 199 200def ensure_dirs(*paths: str) -> None:201    """Create one or more directories if they don't exist."""202    for path in paths:203        os.makedirs(path, exist_ok=True)204        logger.debug(f"Directory ensured: {path}")205 206 207def validate_probe_input(text: str, max_length: int = 2000) -> str:208    """209    Sanitise probe text before sending to model APIs.210    Removes ASCII control characters, limits length, and flags prompt injection patterns.211    """212    if not isinstance(text, str):213        text = str(text or "")214    # Remove control characters except tab and newline215    sanitized = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", text)216    if len(sanitized) > max_length:217        raise ValueError(f"Probe text exceeds maximum allowed length: {len(sanitized)} > {max_length} chars")218 219    injection_patterns = [220        r"ignore (?:all )?previous instructions",221        r"you are now",222        r"disregard (?:all )?prior",223        r"system prompt",224        r"jailbreak",225    ]226    for pattern in injection_patterns:227        if re.search(pattern, sanitized, re.IGNORECASE):228            logger.warning(f"Potential injection pattern flagged in probe input: '{pattern}'")229 230    return sanitized231 232