esmc_generate_embeddings_test.py

from __future__ import annotations
import traceback
import gc
import hashlib
import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import pandas as pd
import torch
from transformers import AutoModel, AutoTokenizer


# ============================================================
# SETTINGS
# ============================================================

INPUT_CSV = Path(
    "/home/dhfn61/masters/embeddings/test_sequences.csv"
)

OUTPUT_DIR = Path(
    "/home/dhfn61/masters/embeddings/ESMC_6B_test"
)

MODEL_NAME = "biohub/ESMC-6B"

ID_COLUMN = "Sequence_id"
SEQUENCE_COLUMN = "Sequence"

# Start conservatively for ESMC-6B. Increase only after checking GPU memory.
BATCH_SIZE = 1

# Saved files use half the storage of float32.
SAVE_DTYPE = torch.float16
SAVE_DTYPE_NAME = "float16"

# Existing embedding files are skipped, making the script resumable.
OVERWRITE_EXISTING = False

# Reduces unnecessary padding within batches.
SORT_UNIQUE_SEQUENCES_BY_LENGTH = True

# Standard amino acids plus commonly used ambiguous/rare one-letter symbols.
# Remove letters here if you want stricter validation.
ALLOWED_RESIDUES = set("ACDEFGHIKLMNPQRSTVWYBXZOUJ")


# ============================================================
# PATHS
# ============================================================

EMBEDDINGS_DIR = OUTPUT_DIR / "embeddings"
MANIFEST_PATH = OUTPUT_DIR / "embedding_manifest.csv"
UNIQUE_INDEX_PATH = OUTPUT_DIR / "unique_embedding_index.csv"
CONFIG_PATH = OUTPUT_DIR / "embedding_config.json"


# ============================================================
# HELPERS
# ============================================================

def utc_now_iso() -> str:
    """Return the current UTC time in an ISO-8601 representation."""
    return datetime.now(timezone.utc).isoformat()


def normalize_sequence(value: Any) -> tuple[str | None, str | None]:
    """
    Clean and validate one peptide sequence.

    Returns
    -------
    cleaned_sequence, error_message
    """
    if pd.isna(value):
        return None, "Missing sequence"

    sequence = re.sub(r"\s+", "", str(value)).upper()

    if not sequence:
        return None, "Empty sequence after whitespace removal"

    invalid_characters = sorted(set(sequence) - ALLOWED_RESIDUES)
    if invalid_characters:
        invalid_text = ", ".join(repr(char) for char in invalid_characters)
        return None, f"Invalid residue character(s): {invalid_text}"

    return sequence, None


def short_sequence_hash(sequence: str, length: int = 20) -> str:
    """Create a stable filename-safe identifier from a sequence."""
    return hashlib.sha256(sequence.encode("utf-8")).hexdigest()[:length]


def atomic_write_csv(dataframe: pd.DataFrame, destination: Path) -> None:
    """Write a CSV through a temporary file to avoid partial outputs."""
    temporary_path = destination.with_suffix(destination.suffix + ".tmp")
    dataframe.to_csv(temporary_path, index=False)
    temporary_path.replace(destination)


def atomic_write_json(data: dict[str, Any], destination: Path) -> None:
    """Write JSON through a temporary file to avoid partial outputs."""
    temporary_path = destination.with_suffix(destination.suffix + ".tmp")
    with temporary_path.open("w", encoding="utf-8") as handle:
        json.dump(data, handle, indent=2, ensure_ascii=False)
    temporary_path.replace(destination)


def atomic_torch_save(data: dict[str, Any], destination: Path) -> None:
    """Write a PyTorch file through a temporary file."""
    temporary_path = destination.with_suffix(destination.suffix + ".tmp")
    torch.save(data, temporary_path)
    temporary_path.replace(destination)


def get_input_device(model: torch.nn.Module) -> torch.device:
    """
    Determine where tokenizer outputs should be placed.

    This follows the same pattern as the official ESMC Hugging Face example,
    while retaining a fallback for model wrappers without model.device.
    """
    try:
        return torch.device(model.device)
    except (AttributeError, TypeError):
        return next(model.parameters()).device


def get_hidden_size(model: torch.nn.Module) -> int | None:
    """Read the model embedding dimension from common config attributes."""
    for attribute in ("hidden_size", "d_model", "embed_dim"):
        value = getattr(model.config, attribute, None)
        if value is not None:
            return int(value)
    return None


def extract_last_hidden_state(output: Any) -> torch.Tensor:
    """
    Extract final residue representations from common Transformers outputs.
    """
    if hasattr(output, "last_hidden_state"):
        hidden_state = output.last_hidden_state
        if hidden_state is not None:
            return hidden_state

    if isinstance(output, dict):
        for key in ("last_hidden_state", "embeddings"):
            hidden_state = output.get(key)
            if hidden_state is not None:
                return hidden_state

    hidden_states = getattr(output, "hidden_states", None)
    if hidden_states:
        return hidden_states[-1]

    raise RuntimeError(
        "The model output did not contain last_hidden_state, embeddings, "
        "or hidden_states."
    )


def is_cuda_out_of_memory(error: BaseException) -> bool:
    """Return True for common CUDA out-of-memory exceptions/messages."""
    if isinstance(error, torch.cuda.OutOfMemoryError):
        return True
    return "out of memory" in str(error).lower()


def clear_memory() -> None:
    """Release Python references and unused CUDA cache."""
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


# ============================================================
# INPUT PREPARATION
# ============================================================

def read_and_prepare_input() -> tuple[pd.DataFrame, list[dict[str, Any]], set[str]]:
    """
    Read the CSV, validate sequences, identify duplicates, and construct one
    embedding job per unique cleaned sequence.
    """
    if not INPUT_CSV.exists():
        raise FileNotFoundError(f"Input file does not exist: {INPUT_CSV}")

    # Read all columns as strings so identifiers such as "000123" keep
    # their leading zeros.
    dataframe = pd.read_csv(INPUT_CSV, dtype="string")

    missing_columns = [
        column
        for column in (ID_COLUMN, SEQUENCE_COLUMN)
        if column not in dataframe.columns
    ]
    if missing_columns:
        raise ValueError(
            f"Missing required column(s): {missing_columns}. "
            f"Available columns: {list(dataframe.columns)}"
        )

    dataframe = dataframe[[ID_COLUMN, SEQUENCE_COLUMN]].copy()
    dataframe.insert(0, "Input_row", range(2, len(dataframe) + 2))

    # Keep the source text visible in the final manifest.
    dataframe[ID_COLUMN] = dataframe[ID_COLUMN].fillna("").astype(str).str.strip()

    normalized = dataframe[SEQUENCE_COLUMN].apply(normalize_sequence)
    dataframe["Cleaned_sequence"] = normalized.map(lambda item: item[0])
    dataframe["Validation_error"] = normalized.map(lambda item: item[1])
    dataframe["Sequence_length"] = dataframe["Cleaned_sequence"].map(
        lambda value: len(value) if isinstance(value, str) else pd.NA
    )

    valid_mask = dataframe["Cleaned_sequence"].notna()
    valid_dataframe = dataframe.loc[valid_mask].copy()

    # Count how often the same cleaned peptide occurs in the source file.
    duplicate_counts = valid_dataframe["Cleaned_sequence"].value_counts()
    dataframe["Duplicate_sequence_count"] = (
        dataframe["Cleaned_sequence"].map(duplicate_counts).fillna(0).astype(int)
    )
    dataframe["Is_duplicate_sequence"] = (
        dataframe["Duplicate_sequence_count"] > 1
    )

    # Detect IDs that point to more than one distinct cleaned sequence.
    nonempty_ids = valid_dataframe[valid_dataframe[ID_COLUMN] != ""]
    sequences_per_id = nonempty_ids.groupby(ID_COLUMN)["Cleaned_sequence"].nunique()
    conflicting_ids = set(sequences_per_id[sequences_per_id > 1].index)
    dataframe["Sequence_id_conflict"] = dataframe[ID_COLUMN].isin(conflicting_ids)

    unique_records: list[dict[str, Any]] = []

    for sequence, group in valid_dataframe.groupby(
        "Cleaned_sequence",
        sort=False,
        dropna=False,
    ):
        sequence_hash = short_sequence_hash(sequence)
        filename = f"sequence_{sequence_hash}.pt"
        relative_path = Path("embeddings") / filename

        sequence_ids = sorted(
            {
                sequence_id
                for sequence_id in group[ID_COLUMN].tolist()
                if sequence_id
            }
        )

        unique_records.append(
            {
                "sequence": sequence,
                "sequence_hash": sequence_hash,
                "sequence_length": len(sequence),
                "sequence_ids": sequence_ids,
                "source_row_count": len(group),
                "filename": filename,
                "relative_path": relative_path.as_posix(),
                "absolute_path": EMBEDDINGS_DIR / filename,
            }
        )

    if SORT_UNIQUE_SEQUENCES_BY_LENGTH:
        unique_records.sort(
            key=lambda record: (
                record["sequence_length"],
                record["sequence_hash"],
            )
        )

    return dataframe, unique_records, conflicting_ids


# ============================================================
# EMBEDDING GENERATION
# ============================================================

def generate_embeddings(
    unique_records: list[dict[str, Any]],
    model: torch.nn.Module,
    tokenizer: Any,
) -> dict[str, dict[str, Any]]:
    """
    Generate and save embeddings.

    Failed batches are recursively split. This lets a smaller batch succeed
    after an out-of-memory error and helps isolate a problematic sequence.
    """
    input_device = get_input_device(model)
    configured_hidden_size = get_hidden_size(model)

    results: dict[str, dict[str, Any]] = {}
    pending_records: list[dict[str, Any]] = []

    for record in unique_records:
        path = record["absolute_path"]

        if path.exists() and not OVERWRITE_EXISTING:
            hidden_size = configured_hidden_size
            token_shape = (
                f"{record['sequence_length']}x{hidden_size}"
                if hidden_size is not None
                else ""
            )
            mean_shape = str(hidden_size) if hidden_size is not None else ""

            results[record["sequence"]] = {
                **record,
                "status": "skipped_existing",
                "error": "",
                "token_embedding_shape": token_shape,
                "mean_embedding_shape": mean_shape,
                "embedding_dimension": hidden_size,
            }
        else:
            pending_records.append(record)

    total_pending = len(pending_records)
    completed_pending = 0

    def mark_failure(record: dict[str, Any], error: BaseException | str) -> None:
        nonlocal completed_pending

        error_message = str(error)
        results[record["sequence"]] = {
            **record,
            "status": "failed",
            "error": error_message,
            "token_embedding_shape": "",
            "mean_embedding_shape": "",
            "embedding_dimension": configured_hidden_size,
        }
        completed_pending += 1
        print(
            f"[{completed_pending}/{total_pending}] FAILED "
            f"{record['sequence_hash']}: {error_message}"
        )

    def process_batch(records: list[dict[str, Any]]) -> None:
        nonlocal completed_pending

        if not records:
            return

        sequences = [record["sequence"] for record in records]

        try:
            encoded = tokenizer(
                sequences,
                return_tensors="pt",
                padding=True,
                return_attention_mask=True,
                return_special_tokens_mask=True,
            )

            # The model does not need special_tokens_mask; it is retained only
            # for selecting amino-acid positions from the output.
            special_tokens_mask = encoded.pop("special_tokens_mask")
            attention_mask_cpu = encoded["attention_mask"].clone()

            model_inputs = {
                key: value.to(input_device)
                for key, value in encoded.items()
            }

            with torch.inference_mode():
                output = model(**model_inputs)

            hidden_state = extract_last_hidden_state(output)

            if hidden_state.ndim != 3:
                raise RuntimeError(
                    "Expected hidden state with shape "
                    "[batch, token_length, embedding_dimension], but received "
                    f"{tuple(hidden_state.shape)}"
                )

            for index, record in enumerate(records):
                try:
                    residue_mask = (
                        attention_mask_cpu[index].bool()
                        & ~special_tokens_mask[index].bool()
                    ).to(hidden_state.device)

                    residue_embeddings = hidden_state[index][residue_mask]

                    expected_length = record["sequence_length"]
                    actual_length = int(residue_embeddings.shape[0])

                    if actual_length != expected_length:
                        raise ValueError(
                            "Residue/token alignment mismatch: "
                            f"sequence length is {expected_length}, but "
                            f"{actual_length} non-special token embeddings "
                            "were produced."
                        )

                    embedding_dimension = int(residue_embeddings.shape[-1])

                    # Mean pooling is calculated in float32 for better numerical
                    # stability and converted only for storage.
                    mean_embedding = (
                        residue_embeddings.detach()
                        .float()
                        .mean(dim=0)
                        .to(device="cpu", dtype=SAVE_DTYPE)
                        .contiguous()
                    )

                    token_embeddings = (
                        residue_embeddings.detach()
                        .to(device="cpu", dtype=SAVE_DTYPE)
                        .contiguous()
                    )

                    payload = {
                        "sequence": record["sequence"],
                        "sequence_hash": record["sequence_hash"],
                        "sequence_ids": record["sequence_ids"],
                        "sequence_length": expected_length,
                        "model_name": MODEL_NAME,
                        "stored_dtype": SAVE_DTYPE_NAME,
                        "special_tokens_removed": True,
                        "token_embeddings": token_embeddings,
                        "mean_embedding": mean_embedding,
                    }

                    atomic_torch_save(payload, record["absolute_path"])

                    results[record["sequence"]] = {
                        **record,
                        "status": "success",
                        "error": "",
                        "token_embedding_shape": (
                            f"{expected_length}x{embedding_dimension}"
                        ),
                        "mean_embedding_shape": str(embedding_dimension),
                        "embedding_dimension": embedding_dimension,
                    }

                    completed_pending += 1
                    print(
                        f"[{completed_pending}/{total_pending}] Saved "
                        f"{record['relative_path']} "
                        f"({expected_length}x{embedding_dimension})"
                    )

                except Exception as item_error:
                    print("\nFULL ITEM TRACEBACK:")
                    traceback.print_exc()
                    mark_failure(record, item_error)

            del output
            del hidden_state
            del model_inputs
            del encoded
            del attention_mask_cpu
            del special_tokens_mask
            clear_memory()

        except Exception as batch_error:
            print("\nFULL BATCH TRACEBACK:")
            traceback.print_exc()
            clear_memory()

            if len(records) > 1:
                reason = (
                    "CUDA out of memory"
                    if is_cuda_out_of_memory(batch_error)
                    else type(batch_error).__name__
                )
                midpoint = len(records) // 2
                print(
                    f"Batch of {len(records)} failed ({reason}); "
                    "retrying as smaller batches."
                )
                process_batch(records[:midpoint])
                process_batch(records[midpoint:])
            else:
                mark_failure(records[0], batch_error)

    for batch_start in range(0, total_pending, BATCH_SIZE):
        batch = pending_records[batch_start : batch_start + BATCH_SIZE]
        process_batch(batch)

    return results


# ============================================================
# OUTPUT TABLES
# ============================================================

def build_and_save_outputs(
    source_dataframe: pd.DataFrame,
    unique_records: list[dict[str, Any]],
    results: dict[str, dict[str, Any]],
    conflicting_ids: set[str],
) -> None:
    """Create the row-level manifest, unique index, and configuration JSON."""
    manifest = source_dataframe.copy()

    def result_value(sequence: Any, key: str, default: Any = "") -> Any:
        if not isinstance(sequence, str):
            return default
        return results.get(sequence, {}).get(key, default)

    invalid_mask = manifest["Cleaned_sequence"].isna()

    manifest["Embedding_file"] = manifest["Cleaned_sequence"].map(
        lambda sequence: result_value(sequence, "relative_path")
    )
    manifest["Embedding_status"] = manifest["Cleaned_sequence"].map(
        lambda sequence: result_value(sequence, "status")
    )
    manifest.loc[invalid_mask, "Embedding_status"] = "invalid"

    manifest["Error"] = manifest["Cleaned_sequence"].map(
        lambda sequence: result_value(sequence, "error")
    )
    manifest.loc[invalid_mask, "Error"] = manifest.loc[
        invalid_mask, "Validation_error"
    ]

    manifest["Token_embedding_shape"] = manifest["Cleaned_sequence"].map(
        lambda sequence: result_value(sequence, "token_embedding_shape")
    )
    manifest["Mean_embedding_shape"] = manifest["Cleaned_sequence"].map(
        lambda sequence: result_value(sequence, "mean_embedding_shape")
    )
    manifest["Embedding_dimension"] = manifest["Cleaned_sequence"].map(
        lambda sequence: result_value(sequence, "embedding_dimension", pd.NA)
    )

    manifest["Model_name"] = MODEL_NAME
    manifest["Stored_dtype"] = SAVE_DTYPE_NAME
    manifest["Special_tokens_removed"] = True

    manifest = manifest[
        [
            "Input_row",
            ID_COLUMN,
            SEQUENCE_COLUMN,
            "Cleaned_sequence",
            "Sequence_length",
            "Embedding_file",
            "Embedding_status",
            "Error",
            "Token_embedding_shape",
            "Mean_embedding_shape",
            "Embedding_dimension",
            "Duplicate_sequence_count",
            "Is_duplicate_sequence",
            "Sequence_id_conflict",
            "Model_name",
            "Stored_dtype",
            "Special_tokens_removed",
        ]
    ]

    unique_index_rows = []

    for record in unique_records:
        result = results.get(record["sequence"], {})

        unique_index_rows.append(
            {
                "Sequence_hash": record["sequence_hash"],
                "Sequence": record["sequence"],
                "Sequence_length": record["sequence_length"],
                "Sequence_ids": "|".join(record["sequence_ids"]),
                "Source_row_count": record["source_row_count"],
                "Embedding_file": record["relative_path"],
                "Embedding_status": result.get("status", "not_processed"),
                "Error": result.get("error", ""),
                "Token_embedding_shape": result.get(
                    "token_embedding_shape", ""
                ),
                "Mean_embedding_shape": result.get(
                    "mean_embedding_shape", ""
                ),
                "Embedding_dimension": result.get(
                    "embedding_dimension", pd.NA
                ),
                "Model_name": MODEL_NAME,
                "Stored_dtype": SAVE_DTYPE_NAME,
            }
        )

    unique_index_columns = [
        "Sequence_hash",
        "Sequence",
        "Sequence_length",
        "Sequence_ids",
        "Source_row_count",
        "Embedding_file",
        "Embedding_status",
        "Error",
        "Token_embedding_shape",
        "Mean_embedding_shape",
        "Embedding_dimension",
        "Model_name",
        "Stored_dtype",
    ]
    unique_index = pd.DataFrame(
        unique_index_rows,
        columns=unique_index_columns,
    )

    atomic_write_csv(manifest, MANIFEST_PATH)
    atomic_write_csv(unique_index, UNIQUE_INDEX_PATH)

    status_counts = manifest["Embedding_status"].value_counts(dropna=False)
    unique_status_counts = unique_index["Embedding_status"].value_counts(
        dropna=False
    )

    config = {
        "created_or_updated_utc": utc_now_iso(),
        "input_csv": str(INPUT_CSV.resolve()),
        "output_directory": str(OUTPUT_DIR.resolve()),
        "model_name": MODEL_NAME,
        "id_column": ID_COLUMN,
        "sequence_column": SEQUENCE_COLUMN,
        "batch_size": BATCH_SIZE,
        "stored_dtype": SAVE_DTYPE_NAME,
        "overwrite_existing": OVERWRITE_EXISTING,
        "sort_unique_sequences_by_length": SORT_UNIQUE_SEQUENCES_BY_LENGTH,
        "special_tokens_removed": True,
        "mean_pooling": "Arithmetic mean over amino-acid token embeddings only",
        "deduplication": "One embedding file per unique cleaned sequence",
        "filename_scheme": "sequence_<first 20 characters of SHA-256>.pt",
        "allowed_residues": "".join(sorted(ALLOWED_RESIDUES)),
        "source_row_count": int(len(manifest)),
        "valid_source_row_count": int(
            (manifest["Embedding_status"] != "invalid").sum()
        ),
        "invalid_source_row_count": int(
            (manifest["Embedding_status"] == "invalid").sum()
        ),
        "unique_valid_sequence_count": int(len(unique_index)),
        "conflicting_sequence_id_count": int(len(conflicting_ids)),
        "row_status_counts": {
            str(key): int(value)
            for key, value in status_counts.items()
        },
        "unique_sequence_status_counts": {
            str(key): int(value)
            for key, value in unique_status_counts.items()
        },
        "embedding_file_contents": {
            "sequence": "Cleaned amino-acid sequence",
            "sequence_hash": "Stable SHA-256-derived sequence identifier",
            "sequence_ids": "All source IDs mapped to this unique sequence",
            "sequence_length": "Number of amino-acid residues",
            "model_name": "Embedding model identifier",
            "stored_dtype": "Tensor storage dtype",
            "special_tokens_removed": True,
            "token_embeddings": "[sequence_length, embedding_dimension]",
            "mean_embedding": "[embedding_dimension]",
        },
    }

    atomic_write_json(config, CONFIG_PATH)


# ============================================================
# MAIN
# ============================================================

def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    EMBEDDINGS_DIR.mkdir(parents=True, exist_ok=True)

    print(f"Reading input: {INPUT_CSV}")
    source_dataframe, unique_records, conflicting_ids = read_and_prepare_input()

    valid_rows = int(source_dataframe["Cleaned_sequence"].notna().sum())
    invalid_rows = int(source_dataframe["Cleaned_sequence"].isna().sum())

    print(f"Source rows: {len(source_dataframe)}")
    print(f"Valid source rows: {valid_rows}")
    print(f"Invalid source rows: {invalid_rows}")
    print(f"Unique valid sequences: {len(unique_records)}")

    if conflicting_ids:
        preview = ", ".join(sorted(conflicting_ids)[:10])
        suffix = " ..." if len(conflicting_ids) > 10 else ""
        print(
            "WARNING: Some Sequence_id values map to multiple different "
            f"sequences ({len(conflicting_ids)} ID(s)): {preview}{suffix}"
        )

    if not unique_records:
        print("No valid sequences were found. Writing reports without loading model.")
        build_and_save_outputs(
            source_dataframe=source_dataframe,
            unique_records=unique_records,
            results={},
            conflicting_ids=conflicting_ids,
        )
        return

    print(f"Loading tokenizer: {MODEL_NAME}")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    print(f"Loading model: {MODEL_NAME}")
    model = AutoModel.from_pretrained(
        MODEL_NAME,
    dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
    ).eval()

    print(f"Model input device: {get_input_device(model)}")
    print(f"Configured hidden size: {get_hidden_size(model)}")

    results = generate_embeddings(
        unique_records=unique_records,
        model=model,
        tokenizer=tokenizer,
    )

    build_and_save_outputs(
        source_dataframe=source_dataframe,
        unique_records=unique_records,
        results=results,
        conflicting_ids=conflicting_ids,
    )

    success_count = sum(
        result["status"] == "success"
        for result in results.values()
    )
    skipped_count = sum(
        result["status"] == "skipped_existing"
        for result in results.values()
    )
    failed_count = sum(
        result["status"] == "failed"
        for result in results.values()
    )

    print("\nFinished.")
    print(f"New embedding files: {success_count}")
    print(f"Existing embedding files skipped: {skipped_count}")
    print(f"Failed unique sequences: {failed_count}")
    print(f"Manifest: {MANIFEST_PATH}")
    print(f"Unique index: {UNIQUE_INDEX_PATH}")
    print(f"Configuration: {CONFIG_PATH}")


if __name__ == "__main__":
    main()