src/produce_test_data_embeddings.py

#!/usr/bin/env python3
"""
Generate protein embeddings using a fine-tuned ProtT5 XL LoRA model.

It loads the base model, attaches LoRA adapters from a local directory or checkpoint, extracts encoder
representations, applies attention-masked mean pooling, and writes a pandas
pickle with protein IDs as index and embedding dimensions as integer columns.
"""

from __future__ import annotations

import argparse
import glob
import json
import math
import re
import sys
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import torch
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Ensure progress logs are written immediately in batch jobs (e.g. SLURM/PBS).
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(line_buffering=True)
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(line_buffering=True)

try:
    from peft import PeftModel
except ImportError as exc:
    raise ImportError(
        "Missing dependency 'peft'. Install with: pip install peft"
    ) from exc

try:
    from safetensors.torch import load_file as load_safetensors
except ImportError:
    load_safetensors = None

ROOT_DIR = Path(__file__).resolve().parent.parent
DEFAULT_VERSION = "ContraMLM_v1_1"
DEFAULT_CSV_NAME = "envhog_test_final_no_leakage.csv"


def fmt_seconds(total_seconds: float) -> str:
    total_seconds = max(0, int(total_seconds))
    hours, rem = divmod(total_seconds, 3600)
    minutes, seconds = divmod(rem, 60)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate embeddings with ProTrans LoRA model.")
    
    # We remove --workdir since paths are now strictly routed to 'data' and 'runs'
    
    parser.add_argument(
        "--version",
        default=DEFAULT_VERSION,
        type=str,
        help=(
            "Model version string used to resolve default paths. "
            "Use 'base' to run inference with the base model (no LoRA adapters)."
        ),
    )
    parser.add_argument(
        "--csv",
        default=str(ROOT_DIR / "data" / DEFAULT_CSV_NAME),
        type=str,
        help="Path to input CSV file with 'id' and 'sequence' columns. Default: data/envhog_test_final_no_leakage.csv",
    )
    parser.add_argument(
        "--output",
        default=None,
        type=str,
        help="Path to output pickle file (.pkl). Default: runs/test_embeddings_protrans_lora_<version>.pkl",
    )
    parser.add_argument(
        "--base-model",
        default="Rostlab/prot_t5_xl_uniref50",
        type=str,
        help="Base Hugging Face model name.",
    )
    parser.add_argument(
        "--adapter-dir",
        default=None,
        type=str,
        help="Path to LoRA adapter directory. Default: <checkpoint-root>/lora_adapters",
    )
    parser.add_argument(
        "--checkpoint-root",
        default=None,
        type=str,
        help="Root path containing checkpoint-* folders. Default: runs/protrans_XL_Full_lora_envhog_<version>",
    )
    parser.add_argument(
        "--batch-size",
        default=2,
        type=int,
        help="Batch size for inference.",
    )
    parser.add_argument(
        "--max-length",
        default=512,
        type=int,
        help="Tokenizer max_length (truncation enabled).",
    )
    parser.add_argument(
        "--max-seqs",
        default=None,
        type=int,
        help="Optional cap on number of sequences (for test runs).",
    )
    parser.add_argument(
        "--progress-every",
        default=50,
        type=int,
        help="Print progress every N batches.",
    )
    parser.add_argument(
        "--save-every-batches",
        default=1000,
        type=int,
        help="Save a checkpoint chunk every N batches. Use 0 to disable.",
    )
    parser.add_argument(
        "--chunk-dir",
        default=None,
        type=str,
        help="Directory for chunk checkpoints. Default: <output_stem>_chunks next to output file.",
    )
    parser.add_argument(
        "--no-final-merge",
        action="store_true",
        help="Do not merge chunks into a single output pickle at the end.",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Overwrite output file if it already exists.",
    )
    parser.add_argument(
        "--save-csv",
        default=None,
        type=str,
        help="Optional path to also export CSV embeddings.",
    )
    parser.add_argument(
        "--report-json",
        default=None,
        type=str,
        help="Path to write adapter validation report JSON. Default: runs/adapter_validation_<version>.json",
    )

    args = parser.parse_args()

    version = args.version
    args.use_base_model = version.lower() == "base"

    default_output_name = f"test_embeddings_protrans_lora_{version}.pkl"
    if args.output is None:
        args.output = str(ROOT_DIR / "runs" / default_output_name)

    default_model_dirname = f"protrans_XL_Full_lora_envhog_{version}"
    if args.checkpoint_root is None:
        args.checkpoint_root = str(ROOT_DIR / "runs" / default_model_dirname)

    if args.adapter_dir is None:
        args.adapter_dir = str(Path(args.checkpoint_root) / "lora_adapters")

    default_report_name = f"adapter_validation_{version}.json"
    if args.report_json is None:
        args.report_json = str(ROOT_DIR / "runs" / default_report_name)

    return args


def read_csv_sequences(csv_path: Path) -> Dict[str, str]:
    df = pd.read_csv(csv_path, usecols=["id", "sequence"])
    if df["id"].duplicated().any():
        n_dups = int(df["id"].duplicated().sum())
        print(f"Warning: {n_dups} duplicate IDs found in CSV; keeping first occurrence.")
        df = df.drop_duplicates(subset="id", keep="first")
    df["sequence"] = (
        df["sequence"]
        .astype(str)
        .str.replace(" ", "", regex=False)
        .str.upper()
        .str.replace("-", "", regex=False)
    )
    return dict(zip(df["id"], df["sequence"]))


def prepare_t5_seq(seq: str) -> str:
    seq = str(seq).replace(" ", "")
    seq = seq.replace("U", "X").replace("Z", "X").replace("O", "X")
    return " ".join(list(seq))


def get_encoder(model):
    if hasattr(model, "encoder"):
        return model.encoder

    get_encoder_fn = getattr(model, "get_encoder", None)
    if callable(get_encoder_fn):
        return get_encoder_fn()

    base_model = getattr(model, "base_model", None)
    if base_model is not None:
        if hasattr(base_model, "encoder"):
            return base_model.encoder
        base_get_encoder_fn = getattr(base_model, "get_encoder", None)
        if callable(base_get_encoder_fn):
            return base_get_encoder_fn()

    inner_model = getattr(model, "model", None)
    if inner_model is not None:
        if hasattr(inner_model, "encoder"):
            return inner_model.encoder
        inner_get_encoder_fn = getattr(inner_model, "get_encoder", None)
        if callable(inner_get_encoder_fn):
            return inner_get_encoder_fn()

    return model


def checkpoint_step(path: str) -> int:
    match = re.search(r"checkpoint-(\d+)", str(path))
    return int(match.group(1)) if match else -1


def load_adapter_state_dict(candidate_dir: Path):
    safe_path = candidate_dir / "adapter_model.safetensors"
    bin_path = candidate_dir / "adapter_model.bin"

    if safe_path.exists() and load_safetensors is not None:
        return load_safetensors(str(safe_path), device="cpu")
    if bin_path.exists():
        return torch.load(str(bin_path), map_location="cpu")
    return None


def is_valid_adapter_state_dict(state_dict) -> bool:
    if state_dict is None:
        return False

    for _, tensor in state_dict.items():
        if not torch.is_tensor(tensor):
            continue
        if torch.isnan(tensor).any().item() or torch.isinf(tensor).any().item():
            return False
    return True


def select_best_adapter_dir(
    adapter_dir: Path,
    checkpoint_root: Path,
) -> Tuple[Path, List[dict]]:
    candidates: List[Path] = []
    reports: List[dict] = []

    if adapter_dir.is_dir():
        candidates.append(adapter_dir)

    checkpoint_dirs = sorted(
        [Path(p) for p in glob.glob(str(checkpoint_root / "checkpoint-*")) if Path(p).is_dir()],
        key=lambda p: checkpoint_step(str(p)),
        reverse=True,
    )
    candidates.extend(checkpoint_dirs)

    unique_candidates: List[Path] = []
    seen = set()
    for candidate in candidates:
        resolved = str(candidate.resolve())
        if resolved not in seen:
            seen.add(resolved)
            unique_candidates.append(candidate)

    if not unique_candidates:
        raise FileNotFoundError("No adapter/checkpoint directories found.")

    for candidate in unique_candidates:
        state_dict = load_adapter_state_dict(candidate)
        valid = is_valid_adapter_state_dict(state_dict)
        reports.append(
            {
                "candidate": str(candidate),
                "valid": bool(valid),
                "has_state_dict": state_dict is not None,
            }
        )

        if valid:
            print(f"Selected adapter directory: {candidate}")
            return candidate, reports

        print(f"Rejected adapter directory: {candidate}")

    raise RuntimeError("No valid adapter/checkpoint directory found (all missing or NaN/Inf).")


def generate_embeddings(
    id2seq: Dict[str, str],
    base_model_name: str,
    adapter_dir: Path,
    checkpoint_root: Path,
    batch_size: int,
    max_length: int,
    max_seqs: Optional[int],
    progress_every: int,
    report_json: Optional[Path],
    output_path: Path,
    save_every_batches: int,
    chunk_dir: Optional[Path],
    final_merge: bool,
    use_base_model: bool = False,
) -> tuple[Optional[pd.DataFrame], List[Path]]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

    print(f"Device: {device}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    print("Loading tokenizer...")
    tokenizer = T5Tokenizer.from_pretrained(base_model_name, do_lower_case=False, legacy=True)

    print("Loading base model...")
    model = T5ForConditionalGeneration.from_pretrained(
        base_model_name,
        torch_dtype=model_dtype,
        low_cpu_mem_usage=True,
    )

    if use_base_model:
        print("Using base model (no LoRA adapters).")
    else:
        print("Selecting best adapter/checkpoint...")
        selected_adapter, reports = select_best_adapter_dir(adapter_dir, checkpoint_root)

        if report_json is not None:
            report_json.parent.mkdir(parents=True, exist_ok=True)
            with report_json.open("w", encoding="utf-8") as f:
                json.dump(reports, f, indent=2)
            print(f"Adapter validation report saved: {report_json}")

        print("Attaching LoRA adapters...")
        model = PeftModel.from_pretrained(model, str(selected_adapter))

        merge_fn = getattr(model, "merge_and_unload", None)
        if callable(merge_fn):
            model = merge_fn()
            print("Merged LoRA adapters into base weights.")

    model = model.to(device).eval()
    encoder = get_encoder(model).to(device).eval()

    all_ids = list(id2seq.keys())
    if max_seqs is not None:
        all_ids = all_ids[:max_seqs]

    if batch_size <= 0:
        raise ValueError("--batch-size must be > 0")

    total_sequences = len(all_ids)
    total_batches = math.ceil(total_sequences / batch_size)

    lengths = [len(id2seq[pid]) for pid in all_ids]
    print(
        "Input length stats | "
        f"mean={np.mean(lengths):.1f}, median={np.median(lengths):.1f}, "
        f"p95={np.percentile(lengths, 95):.1f}, max={np.max(lengths)}"
    )

    print(f"Generating embeddings for {total_sequences} proteins in ~{total_batches} batches...")
    chunk_arrays: List[np.ndarray] = []
    chunk_ids: List[str] = []
    chunk_paths: List[Path] = []

    if chunk_dir is None:
        chunk_dir = output_path.parent / f"{output_path.stem}_chunks"
    chunk_dir.mkdir(parents=True, exist_ok=True)
    print(f"Checkpoint chunk directory: {chunk_dir}")

    start = time.time()
    processed = 0

    def flush_chunk() -> None:
        if not chunk_arrays:
            return

        chunk_index = len(chunk_paths)
        chunk_path = chunk_dir / f"chunk_{chunk_index:06d}.pkl"
        chunk_emb = np.vstack(chunk_arrays)
        chunk_df = pd.DataFrame(chunk_emb, index=chunk_ids)
        chunk_df.columns = list(range(chunk_df.shape[1]))
        chunk_df.to_pickle(chunk_path)

        chunk_paths.append(chunk_path)
        print(
            f"[checkpoint] saved {chunk_path.name} with {chunk_df.shape[0]} proteins "
            f"(total processed={processed}/{total_sequences})"
        )

        chunk_arrays.clear()
        chunk_ids.clear()

    with torch.no_grad():
        for i in range(0, len(all_ids), batch_size):
            batch_ids = all_ids[i : i + batch_size]
            batch_seqs = [prepare_t5_seq(id2seq[pid]) for pid in batch_ids]

            inputs = tokenizer(
                batch_seqs,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=max_length,
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}

            outputs = encoder(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
            )
            hidden = outputs.last_hidden_state
            mask = inputs["attention_mask"].unsqueeze(-1).to(hidden.dtype)
            pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
            chunk_arrays.append(pooled.float().cpu().numpy())
            chunk_ids.extend(batch_ids)
            processed += len(batch_ids)

            batch_idx = i // batch_size
            if (
                progress_every > 0
                and ((batch_idx + 1) % progress_every == 0 or processed == total_sequences or batch_idx == 0)
            ):
                elapsed = time.time() - start
                seq_per_sec = processed / elapsed if elapsed > 0 else 0.0
                remaining = total_sequences - processed
                eta_seconds = (remaining / seq_per_sec) if seq_per_sec > 0 else float("inf")

                msg = (
                    f"[{batch_idx + 1}/{total_batches}] "
                    f"processed={processed}/{total_sequences} ({processed/total_sequences:.1%}) | "
                    f"throughput={seq_per_sec:.2f} seq/s | "
                    f"elapsed={fmt_seconds(elapsed)} | "
                    f"eta={fmt_seconds(eta_seconds) if math.isfinite(eta_seconds) else 'inf'}"
                )

                if torch.cuda.is_available():
                    alloc_gb = torch.cuda.memory_allocated() / (1024 ** 3)
                    reserved_gb = torch.cuda.memory_reserved() / (1024 ** 3)
                    msg += f" | gpu_mem={alloc_gb:.2f}/{reserved_gb:.2f} GB"

                print(msg)

            if save_every_batches > 0 and ((batch_idx + 1) % save_every_batches == 0):
                flush_chunk()

    flush_chunk()

    if not chunk_paths:
        raise RuntimeError("No embeddings were generated.")

    elapsed = time.time() - start
    print(f"Embedding generation completed in {elapsed:.1f}s")
    print(f"Average throughput: {processed / elapsed:.2f} seq/s")
    print(f"Saved chunk files: {len(chunk_paths)}")

    if not final_merge:
        print("Skipping final merge (--no-final-merge set).")
        return None, chunk_paths

    print("Merging chunk files into final DataFrame...")
    frames: List[pd.DataFrame] = []
    for i, chunk_path in enumerate(chunk_paths, start=1):
        frames.append(pd.read_pickle(chunk_path))
        if i == 1 or i % 50 == 0 or i == len(chunk_paths):
            print(f"[merge] loaded {i}/{len(chunk_paths)} chunks")

    emb_df = pd.concat(frames, axis=0)
    print(f"Merged output shape: {emb_df.shape}")
    return emb_df, chunk_paths


def main() -> None:
    args = parse_args()

    csv_path = Path(args.csv)
    output_path = Path(args.output)
    adapter_dir = Path(args.adapter_dir)
    checkpoint_root = Path(args.checkpoint_root)
    report_json = Path(args.report_json) if args.report_json else None
    chunk_dir = Path(args.chunk_dir) if args.chunk_dir else None

    print("Resolved configuration:")
    print(f"  workdir: {Path(args.workdir)}")
    print(f"  version: {args.version}")
    print(f"  use_base_model: {args.use_base_model}")
    print(f"  csv: {csv_path}")
    print(f"  output: {output_path}")
    print(f"  adapter_dir: {adapter_dir}")
    print(f"  checkpoint_root: {checkpoint_root}")
    print(f"  report_json: {report_json}")
    print(f"  save_every_batches: {args.save_every_batches}")
    print(f"  chunk_dir: {chunk_dir}")
    print(f"  no_final_merge: {args.no_final_merge}")

    if not csv_path.exists():
        raise FileNotFoundError(f"CSV file not found: {csv_path}")

    if output_path.exists() and not args.overwrite:
        raise FileExistsError(
            f"Output already exists: {output_path}. Use --overwrite to replace it."
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)

    print("Reading CSV...")
    id2seq = read_csv_sequences(csv_path)
    print(f"Loaded {len(id2seq)} protein sequences.")

    emb_df, chunk_paths = generate_embeddings(
        id2seq=id2seq,
        base_model_name=args.base_model,
        adapter_dir=adapter_dir,
        checkpoint_root=checkpoint_root,
        batch_size=args.batch_size,
        max_length=args.max_length,
        max_seqs=args.max_seqs,
        progress_every=args.progress_every,
        report_json=report_json,
        output_path=output_path,
        save_every_batches=args.save_every_batches,
        chunk_dir=chunk_dir,
        final_merge=not args.no_final_merge,
        use_base_model=args.use_base_model,
    )

    if emb_df is not None:
        emb_df.to_pickle(output_path)
        print(f"Saved pickle embeddings: {output_path}")
    else:
        print("Final pickle not written because final merge was skipped.")

    if args.save_csv:
        save_csv_path = Path(args.save_csv)
        save_csv_path.parent.mkdir(parents=True, exist_ok=True)
        if emb_df is not None:
            emb_df.to_csv(save_csv_path)
            print(f"Saved csv embeddings: {save_csv_path}")
        else:
            # Stream chunk files to csv when final merge is skipped.
            header_written = False
            for chunk_path in chunk_paths:
                chunk_df = pd.read_pickle(chunk_path)
                chunk_df.to_csv(save_csv_path, mode="a" if header_written else "w", header=not header_written)
                header_written = True
            print(f"Saved csv embeddings from chunks: {save_csv_path}")


if __name__ == "__main__":
    main()