Upload folder using huggingface_hub

README.md

@@ -1,3 +1,112 @@
 ---
-license: mit
+pipeline_tag: feature-extraction
+tags:
+- biology
+- genomics
+- bioinformatics
+- protein-language-model
+- pytorch
+- peft
+- lora
+license: apache-2.0
+language:
+- en
 ---
+
+# ContraMLM: Contrastive Learning for Phage Protein Representations
+
+ContraMLM is a fine-tuned version of the `Rostlab/prot_t5_xl_uniref50` protein language model, trained with Low-Rank Adaptation (LoRA) using a hybrid objective that combines standard Masked Language Modeling (MLM) with a contrastive loss.
+
+The model is built to improve the embedding space for bacteriophage proteins, clustering them by functional group and PHROG family in a zero-shot setting.
+
+## Intended Use
+
+ContraMLM is intended for researchers in computational biology and virology who need function-aware embeddings for phage protein sequences.
+
+**Primary use cases:**
+- **Zero-shot functional retrieval:** querying unknown phage proteins against a database of known PHROG families using cosine similarity or L2 distance.
+- **Embedding extraction:** generating dense representations of sequences for downstream tasks such as functional annotation or clustering.
+
+## How to Use
+
+The model relies on the Hugging Face `transformers` and `peft` libraries.
+
+```python
+import torch
+from transformers import T5Tokenizer, T5ForConditionalGeneration
+from peft import PeftModel
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+base_model_name = "Rostlab/prot_t5_xl_uniref50"
+adapter_dir = "Amin-Saeidi/PhageContraMLM"  # update to your exact repo ID
+
+# 1. Load tokenizer and base model
+tokenizer = T5Tokenizer.from_pretrained(base_model_name, do_lower_case=False)
+model = T5ForConditionalGeneration.from_pretrained(base_model_name, torch_dtype=torch.float16)
+
+# 2. Attach LoRA adapters and merge
+model = PeftModel.from_pretrained(model, adapter_dir)
+model = model.merge_and_unload().to(device).eval()
+
+# 3. Prepare sequence (space-separated, rare amino acids replaced)
+seq = "M A K K L K I L L L A A S L V S L S P S V F A"
+inputs = tokenizer(seq, return_tensors="pt").to(device)
+
+# 4. Extract mean-pooled embeddings
+with torch.no_grad():
+    outputs = model.encoder(**inputs)
+    hidden = outputs.last_hidden_state
+    mask = inputs.attention_mask.unsqueeze(-1).to(hidden.dtype)
+    pooled_embedding = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
+
+print(pooled_embedding.shape)
+```
+
+## Training Data and Process
+
+- **Data:** `envhog_phrog2` dataset.
+- **Architecture:** ProtT5-XL (encoder-decoder).
+- **Fine-tuning method:** LoRA (rank=32, alpha=64, target modules: q, k, v, o).
+- **Loss function:** convex combination of curriculum-aware MLM loss (80%) and temperature-scaled contrastive loss (20%).
+- **Positive sampling:** the contrastive adjacency matrix is built on-the-fly from a known VISEQ pair graph, pulling positive pairs that share identical VISEQs or cross-VISEQ structural similarities.
+
+## Repository Structure & Scripts
+
+This repository contains everything needed to reproduce training and evaluation, organized around `src/`, `data/`, and `runs/` directories.
+
+**`src/train.py`**
+Main training loop. Implements a custom `PairGraphCollator` that samples positive protein pairs on-the-fly and builds a dynamic adjacency matrix for the contrastive loss, combined with a curriculum-aware MLM objective.
+
+**`src/produce_test_data_embeddings.py`**
+High-throughput script for generating mean-pooled encoder embeddings. Loads the base ProtT5 model, attaches the best LoRA adapters from your checkpoints, and processes raw FASTA/CSV sequences in batches, saving results as `.pkl` and `.csv` files.
+
+**`src/eval_EmbeddingSpace.py`**
+Generates publication-quality plots analyzing the embedding space:
+- t-SNE grids colored by PhrogCat category
+- Pairwise L2 and cosine scatter plots comparing the base model against the fine-tuned ContraMLM model
+
+**`src/eval_PhrogRetrieval.py`**
+Zero-shot functional retrieval benchmarking using `hnswlib` (Hierarchical Navigable Small World graphs):
+- Precision@k (k = 5, 10, 50)
+- Per-family and per-size-bin (rare, medium, common) statistics
+- Functional group confusion matrices and Seaborn clustermaps
+
+## Dependencies
+
+```
+torch==2.6.0
+transformers==4.37.2
+peft==0.10.0
+pandas==3.0.1
+numpy==2.3.5
+matplotlib==3.10.8
+seaborn==0.13.2
+scikit-learn==1.8.0
+hnswlib==0.8.0
+safetensors==0.7.0
+sentencepiece==0.2.0
+```
+
+## Acknowledgments
+
+This work was conducted by Amin SaeidiKelishami during an internship at the Laboratoire Microorganismes: Génome et Environnement (LMGE), under the supervision of Professor Clovis Galiez and Professor Francois ENAULT.

pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "phage-contramlm"
+version = "1.1.0"
+description = "A fine-tuned ProtT5 model using Contrastive + MLM loss for phage protein sequence representation."
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "torch==2.6.0",
+    "transformers==4.37.2",
+    "peft==0.10.0",
+    "pandas==3.0.1",
+    "numpy==2.3.5",
+    "matplotlib==3.10.8",
+    "seaborn==0.13.2",
+    "scikit-learn==1.8.0",
+    "hnswlib==0.8.0",
+    "safetensors==0.7.0",
+    "sentencepiece==0.2.0"
+]

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.6.0
+transformers==4.37.2
+peft==0.10.0
+pandas==3.0.1
+numpy==2.3.5
+matplotlib==3.10.8
+seaborn==0.13.2
+scikit-learn==1.8.0
+hnswlib==0.8.0
+safetensors==0.7.0
+sentencepiece==0.2.0

requirements_conda.txt

@@ -0,0 +1,133 @@
+# This file may be used to create an environment using:
+# $ conda create --name <env> --file <this file>
+# platform: linux-64
+_libgcc_mutex=0.1=main
+_openmp_mutex=5.1=1_gnu
+accelerate=0.22.0=pypi_0
+aiohappyeyeballs=2.6.1=pypi_0
+aiohttp=3.13.3=pypi_0
+aiosignal=1.4.0=pypi_0
+annotated-doc=0.0.4=pypi_0
+anyio=4.12.1=pypi_0
+attrs=25.4.0=pypi_0
+biopython=1.86=pypi_0
+bzip2=1.0.8=h5eee18b_6
+ca-certificates=2026.2.25=hbd8a1cb_0
+certifi=2026.2.25=pypi_0
+charset-normalizer=3.4.5=pypi_0
+click=8.3.1=pypi_0
+contourpy=1.3.3=pypi_0
+cycler=0.12.1=pypi_0
+datasets=4.7.0=pypi_0
+dill=0.4.0=pypi_0
+et-xmlfile=2.0.0=pypi_0
+expat=2.7.4=h7354ed3_0
+filelock=3.20.0=pypi_0
+fonttools=4.62.0=pypi_0
+frozenlist=1.8.0=pypi_0
+fsspec=2025.12.0=pypi_0
+h11=0.16.0=pypi_0
+hf-xet=1.4.0=pypi_0
+hnswlib=0.8.0=py312hf890105_1
+httpcore=1.0.9=pypi_0
+httpx=0.28.1=pypi_0
+huggingface-hub=0.36.2=pypi_0
+idna=3.11=pypi_0
+jinja2=3.1.6=pypi_0
+joblib=1.5.3=pypi_0
+kiwisolver=1.5.0=pypi_0
+ld_impl_linux-64=2.44=h9e0c5a2_3
+libblas=3.11.0=6_h4a7cf45_openblas
+libcblas=3.11.0=6_h0358290_openblas
+libexpat=2.7.4=h7354ed3_0
+libffi=3.4.4=h6a678d5_1
+libgcc=15.2.0=h69a1729_7
+libgcc-ng=15.2.0=h166f726_7
+libgfortran=15.2.0=h69a702a_18
+libgfortran5=15.2.0=h68bc16d_18
+libgomp=15.2.0=h4751f2c_7
+liblapack=3.11.0=6_h47877c9_openblas
+libnsl=2.0.0=h5eee18b_0
+libopenblas=0.3.32=pthreads_h94d23a6_0
+libstdcxx=15.2.0=h39759b7_7
+libstdcxx-ng=15.2.0=hc03a8fd_7
+libuuid=1.41.5=h5eee18b_0
+libxcb=1.17.0=h9b100fa_0
+libzlib=1.3.1=hb25bd0a_0
+markdown-it-py=4.0.0=pypi_0
+markupsafe=3.0.2=pypi_0
+matplotlib=3.10.8=pypi_0
+mdurl=0.1.2=pypi_0
+mpmath=1.3.0=pypi_0
+multidict=6.7.1=pypi_0
+multiprocess=0.70.18=pypi_0
+ncurses=6.5=h7934f7d_0
+networkx=3.6.1=pypi_0
+numpy=2.3.5=pypi_0
+nvidia-cublas-cu12=12.4.5.8=pypi_0
+nvidia-cuda-cupti-cu12=12.4.127=pypi_0
+nvidia-cuda-nvrtc-cu12=12.4.127=pypi_0
+nvidia-cuda-runtime-cu12=12.4.127=pypi_0
+nvidia-cudnn-cu12=9.1.0.70=pypi_0
+nvidia-cufft-cu12=11.2.1.3=pypi_0
+nvidia-curand-cu12=10.3.5.147=pypi_0
+nvidia-cusolver-cu12=11.6.1.9=pypi_0
+nvidia-cusparse-cu12=12.3.1.170=pypi_0
+nvidia-cusparselt-cu12=0.6.2=pypi_0
+nvidia-nccl-cu12=2.21.5=pypi_0
+nvidia-nvjitlink-cu12=12.4.127=pypi_0
+nvidia-nvtx-cu12=12.4.127=pypi_0
+openpyxl=3.1.5=pypi_0
+openssl=3.6.1=h35e630c_1
+packaging=25.0=py312h06a4308_1
+pandas=3.0.1=pypi_0
+peft=0.10.0=pypi_0
+pillow=12.0.0=pypi_0
+pip=26.0.1=pyhc872135_0
+propcache=0.4.1=pypi_0
+protobuf=3.20.3=pypi_0
+psutil=7.2.2=pypi_0
+pthread-stubs=0.3=h0ce48e5_1
+pyarrow=23.0.1=pypi_0
+pygments=2.19.2=pypi_0
+pyparsing=3.3.2=pypi_0
+python=3.12.12=hd17a9e1_1
+python-dateutil=2.9.0.post0=pypi_0
+python_abi=3.12=3_cp312
+pyyaml=6.0.3=pypi_0
+readline=8.3=hc2a1206_0
+regex=2026.2.28=pypi_0
+requests=2.32.5=pypi_0
+rich=14.3.3=pypi_0
+safetensors=0.7.0=pypi_0
+scikit-learn=1.8.0=pypi_0
+scipy=1.17.1=pypi_0
+seaborn=0.13.2=pypi_0
+sentencepiece=0.2.0=pypi_0
+setuptools=80.10.2=py312h06a4308_0
+shellingham=1.5.4=pypi_0
+six=1.17.0=pypi_0
+sqlite=3.51.1=h3e8d24a_1
+sympy=1.13.1=pypi_0
+threadpoolctl=3.6.0=pypi_0
+tk=8.6.15=h54e0aa7_0
+tokenizers=0.15.2=pypi_0
+torch=2.6.0+cu124=pypi_0
+torchaudio=2.6.0+cu124=pypi_0
+torchvision=0.21.0+cu124=pypi_0
+tqdm=4.67.3=pypi_0
+transformers=4.37.2=pypi_0
+triton=3.2.0=pypi_0
+typer=0.24.1=pypi_0
+typing-extensions=4.15.0=pypi_0
+tzdata=2026a=he532380_0
+urllib3=2.6.3=pypi_0
+wheel=0.46.3=py312h06a4308_0
+xorg-libx11=1.8.12=h9b100fa_1
+xorg-libxau=1.0.12=h9b100fa_0
+xorg-libxdmcp=1.1.5=h9b100fa_0
+xorg-xorgproto=2024.1=h5eee18b_1
+xxhash=3.6.0=pypi_0
+xz=5.8.2=h448239c_0
+yarl=1.23.0=pypi_0
+zlib=1.3.1=hb25bd0a_0

runs/test_embeddings_protrans_lora_ContraMLM_v1_1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c84c5728b5bc876b571a7e8dddb75c4cbaa7d0242798f393968e623ac0a6717
+size 872815878

runs/test_embeddings_protrans_lora_base.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70723f01c1ab0e4164beb6065e00e90611acd7c1aba0de5c52e26c9b9ee99a30
+size 872815878

src/produce_test_data_embeddings.py

@@ -0,0 +1,570 @@
+#!/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()