DeepHealthExpo/encode_exposure_cache.py

"""Encode cached exposure windows once for embedding-only model training."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

from backbones import TimesNetExposureEncoder
from dataset import ExposureCache
from train_util import configure_torch_for_training, resolve_device


class ExposureEncodingDataset(Dataset):
    def __init__(self, cache: ExposureCache):
        self.cache = cache

    def __len__(self) -> int:
        return len(self.cache.daily)

    def __getitem__(self, index: int) -> tuple[int, torch.Tensor, torch.Tensor]:
        return (
            index,
            torch.from_numpy(
                np.array(self.cache.daily[index], dtype=np.float32, copy=True)
            ),
            torch.from_numpy(
                np.array(self.cache.monthly[index], dtype=np.float32, copy=True)
            ),
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Precompute exposure embeddings from an autoencoder checkpoint"
    )
    parser.add_argument("--exposure_cache_dir", required=True)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--output_file", default=None)
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument("--device", default="cuda")
    parser.add_argument(
        "--output_dtype", choices=["float16", "float32"], default="float16"
    )
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args()
    if args.batch_size <= 0:
        parser.error("--batch_size must be positive")
    return args


def main() -> None:
    args = parse_args()
    device = resolve_device(args.device)
    configure_torch_for_training(device)
    cache_dir = Path(args.exposure_cache_dir)
    output_path = (
        Path(args.output_file)
        if args.output_file
        else cache_dir / "exposure_embeddings.npy"
    )
    if output_path.exists() and not args.overwrite:
        raise FileExistsError(f"{output_path} exists; pass --overwrite")

    try:
        checkpoint_data = torch.load(
            args.checkpoint, map_location="cpu", weights_only=False
        )
    except TypeError:
        checkpoint_data = torch.load(args.checkpoint, map_location="cpu")
    model_cfg = checkpoint_data["model_config"]
    normalization = checkpoint_data["normalization"]
    encoder = TimesNetExposureEncoder(
        n_embd=int(model_cfg["n_embd"]),
        d_model=model_cfg["d_model"],
        n_layers=int(model_cfg["n_layers"]),
        top_k=int(model_cfg["top_k"]),
        n_backbone_blocks=int(model_cfg["n_backbone_blocks"]),
        backbone_kernel_size=int(model_cfg["backbone_kernel_size"]),
        backbone_expansion=float(model_cfg["backbone_expansion"]),
        dropout=float(model_cfg["dropout"]),
        use_gate=True,
    )
    encoder.load_state_dict(checkpoint_data["encoder_state_dict"], strict=True)
    encoder.to(device).eval()

    cache = ExposureCache(cache_dir)
    loader_kwargs = {
        "batch_size": args.batch_size,
        "shuffle": False,
        "num_workers": args.num_workers,
        "pin_memory": device.type == "cuda",
        "persistent_workers": args.num_workers > 0,
    }
    loader = DataLoader(ExposureEncodingDataset(cache), **loader_kwargs)
    output_dtype = np.float16 if args.output_dtype == "float16" else np.float32
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output = np.lib.format.open_memmap(
        output_path,
        mode="w+",
        dtype=output_dtype,
        shape=(len(cache.daily), int(model_cfg["n_embd"])),
    )
    stats = {
        key: torch.as_tensor(normalization[key], device=device).view(1, 1, -1)
        for key in ("daily_mean", "daily_std", "monthly_mean", "monthly_std")
    }

    with torch.inference_mode():
        for indices, daily, monthly in tqdm(loader, desc="Encoding exposure"):
            daily = daily.to(device, non_blocking=True)
            monthly = monthly.to(device, non_blocking=True)
            daily_mask = torch.isfinite(daily)
            monthly_mask = torch.isfinite(monthly)
            has_observation = (
                daily_mask.flatten(1).any(dim=1)
                | monthly_mask.flatten(1).any(dim=1)
            )
            daily = (
                torch.nan_to_num(daily) - stats["daily_mean"]
            ) / stats["daily_std"]
            monthly = (
                torch.nan_to_num(monthly) - stats["monthly_mean"]
            ) / stats["monthly_std"]
            daily = daily * daily_mask
            monthly = monthly * monthly_mask
            encoded = encoder(daily, monthly, daily_mask, monthly_mask)
            encoded = encoded * has_observation.unsqueeze(1)
            output[indices.numpy()] = encoded.float().cpu().numpy().astype(
                output_dtype, copy=False
            )
    output.flush()
    metadata = {
        "checkpoint": str(Path(args.checkpoint).resolve()),
        "output_file": str(output_path.resolve()),
        "rows": len(cache.daily),
        "embedding_dim": int(model_cfg["n_embd"]),
        "dtype": args.output_dtype,
    }
    output_path.with_suffix(".json").write_text(
        json.dumps(metadata, indent=2), encoding="utf-8"
    )
    print(f"Wrote {len(cache.daily):,} embeddings to {output_path}")


if __name__ == "__main__":
    main()