Upload modeling_esm_plusplus.py with huggingface_hub

modeling_esm_plusplus.py

@@ -1,4 +1,1015 @@
 from __future__ import annotations
+
+import torch
+import torch._inductor.config as inductor_config
+import torch._dynamo as dynamo
+
+# Enable TensorFloat32 tensor cores for float32 matmul (Ampere+ GPUs)
+# Provides significant speedup with minimal precision loss
+torch.set_float32_matmul_precision('high')
+
+# Enable TF32 for matrix multiplications and cuDNN operations
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+# Enable cuDNN autotuner - finds fastest algorithms for your hardware
+# Best when input sizes are consistent; may slow down first iterations
+torch.backends.cudnn.benchmark = True
+
+# Deterministic operations off for speed (set True if reproducibility needed)
+torch.backends.cudnn.deterministic = False
+inductor_config.max_autotune_gemm_backends = "ATEN,CUTLASS,FBGEMM"    
+
+dynamo.config.capture_scalar_outputs = True
+torch._dynamo.config.recompile_limit = 16
+
+import io
+import os
+import queue
+import sqlite3
+import struct
+import threading
+import time
+
+import networkx as nx
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset as TorchDataset
+from transformers import PreTrainedTokenizerBase
+
+
+# Compact blob serialization constants
+# Canonical source: core/embed/blob.py. Keep in sync with protify/utils.py.
+_COMPACT_VERSION = 0x01
+_DTYPE_TO_CODE = {torch.float16: 0, torch.bfloat16: 1, torch.float32: 2}
+_CODE_TO_DTYPE = {0: torch.float16, 1: torch.bfloat16, 2: torch.float32}
+_CODE_TO_NP_DTYPE = {0: np.float16, 1: np.float16, 2: np.float32}
+
+
+def tensor_to_embedding_blob(tensor: torch.Tensor) -> bytes:
+    """Serialize a tensor to compact binary format for SQLite blob storage.
+
+    Format: [version:1][dtype_code:1][ndim:4][shape:4*ndim][raw_bytes]
+    bfloat16 tensors are stored as float16 bytes (numpy lacks bfloat16)
+    but tagged with dtype_code=1 so they can be cast back on read.
+    Falls back to torch.save for unsupported dtypes.
+    """
+    t = tensor.cpu()
+    if t.dtype not in _DTYPE_TO_CODE:
+        buffer = io.BytesIO()
+        torch.save(t, buffer)
+        return buffer.getvalue()
+    dtype_code = _DTYPE_TO_CODE[t.dtype]
+
+    if t.dtype == torch.bfloat16:
+        raw = t.half().numpy().tobytes()
+    else:
+        raw = t.numpy().tobytes()
+
+    shape = t.shape
+    header = struct.pack(f'<BBi{len(shape)}i', _COMPACT_VERSION, dtype_code, len(shape), *shape)
+    return header + raw
+
+
+def _compact_header(dtype: torch.dtype, shape: tuple) -> bytes:
+    """Build just the compact header for a given dtype and shape."""
+    dtype_code = _DTYPE_TO_CODE[dtype]
+    return struct.pack(f'<BBi{len(shape)}i', _COMPACT_VERSION, dtype_code, len(shape), *shape)
+
+
+def batch_tensor_to_blobs(batch: torch.Tensor) -> List[bytes]:
+    """Serialize a batch of same-shape tensors to compact blobs (fast path for vectors).
+
+    Builds the header once and slices raw bytes per row. Much faster than
+    per-row tensor_to_embedding_blob calls for uniform-shape batches.
+    """
+    assert batch.ndim >= 2, f"Expected batch with >= 2 dims, got {batch.ndim}"
+    t = batch.cpu()
+    store_dtype = t.dtype
+    if t.dtype not in _DTYPE_TO_CODE:
+        return [tensor_to_embedding_blob(t[i]) for i in range(t.shape[0])]
+
+    if t.dtype == torch.bfloat16:
+        arr = t.half().numpy()
+        store_dtype = torch.bfloat16
+    else:
+        arr = t.numpy()
+
+    row_shape = tuple(t.shape[1:])
+    header = _compact_header(store_dtype, row_shape)
+    raw = arr.tobytes()
+    stride = len(raw) // t.shape[0]
+    return [header + raw[i * stride:(i + 1) * stride] for i in range(t.shape[0])]
+
+
+def embedding_blob_to_tensor(blob: bytes, fallback_shape: Optional[Tuple[int, ...]] = None) -> torch.Tensor:
+    """Deserialize a blob back to a tensor. Auto-detects compact vs legacy formats."""
+    if len(blob) >= 6 and blob[0] == _COMPACT_VERSION:
+        dtype_code = blob[1]
+        ndim = struct.unpack_from('<i', blob, 2)[0]
+        shape = struct.unpack_from(f'<{ndim}i', blob, 6)
+        data_offset = 6 + 4 * ndim
+        np_dtype = _CODE_TO_NP_DTYPE[dtype_code]
+        arr = np.frombuffer(blob, dtype=np_dtype, offset=data_offset).copy().reshape(shape)
+        t = torch.from_numpy(arr)
+        target_dtype = _CODE_TO_DTYPE[dtype_code]
+        if target_dtype != t.dtype:
+            t = t.to(target_dtype)
+        return t
+
+    # Fallback: try torch.load (pickle format)
+    try:
+        buffer = io.BytesIO(blob)
+        return torch.load(buffer, map_location='cpu', weights_only=True)
+    except Exception:
+        pass
+
+    # Legacy fallback: raw float32 bytes with caller-supplied shape
+    assert fallback_shape is not None, "Cannot deserialize blob: unknown format and no fallback_shape provided."
+    arr = np.frombuffer(blob, dtype=np.float32).copy().reshape(fallback_shape)
+    return torch.from_numpy(arr)
+
+
+def maybe_compile(model: torch.nn.Module, dynamic: bool = False) -> torch.nn.Module:
+    """Compile model with torch.compile if possible.
+
+    Skips compilation when dynamic=True (padding='longest') because
+    flex attention's create_block_mask is incompatible with dynamic shapes
+    under torch.compile, causing CUDA illegal memory access.
+    """
+    if dynamic:
+        print("Skipping torch.compile (dynamic shapes + flex attention incompatible)")
+        return model
+    try:
+        model = torch.compile(model)
+        print("Model compiled")
+    except Exception as e:
+        print(f"Skipping torch.compile: {e}")
+    return model
+
+
+def build_collator(
+    tokenizer: PreTrainedTokenizerBase,
+    padding: str = 'max_length',
+    max_length: int = 512,
+) -> Callable[[List[str]], Dict[str, torch.Tensor]]:
+    def _collate_fn(sequences: List[str]) -> Dict[str, torch.Tensor]:
+        kwargs: Dict[str, Any] = dict(
+            return_tensors="pt", padding=padding, truncation=True, max_length=max_length,
+        )
+        if padding != 'max_length':
+            kwargs['pad_to_multiple_of'] = 8
+        return tokenizer(sequences, **kwargs)
+    return _collate_fn
+
+
+def _make_embedding_progress(
+    dataloader: DataLoader,
+    padding: str,
+    n_warmup: int = 3,
+    n_calibration: int = 5,
+) -> Iterator[Tuple[int, Any]]:
+    """Progress-bar wrapper for embedding loops. Drop-in replacement for enumerate(dataloader).
+
+    When padding='max_length', all batches have uniform cost so plain tqdm works.
+    When padding='longest' (sorted longest-first), batch times vary dramatically.
+    In that case: yield warmup batches first (compiler warmup + OOM check on longest
+    sequences), then time mid-length calibration batches to estimate total ETA.
+
+    Keep in sync with protify/embedder.py and core/atlas/precomputed.py.
+    """
+    total = len(dataloader)
+    if padding == 'max_length' or total <= n_warmup + n_calibration:
+        for i, batch in tqdm(enumerate(dataloader), total=total, desc='Embedding batches'):
+            yield i, batch
+        return
+
+    dl_iter = iter(dataloader)
+
+    # Phase 1: warmup on longest batches (first n_warmup, since sorted longest-first)
+    warmup_bar = tqdm(range(n_warmup), desc='Warmup (longest batches)', leave=False)
+    for i in warmup_bar:
+        batch = next(dl_iter)
+        yield i, batch
+    warmup_bar.close()
+
+    # Phase 2: skip to middle of dataset for calibration timing
+    # We need to yield all intermediate batches too (they contain real data)
+    mid_start = total // 2
+    intermediate_bar = tqdm(
+        range(n_warmup, mid_start), desc='Embedding batches', leave=False,
+    )
+    for i in intermediate_bar:
+        batch = next(dl_iter)
+        yield i, batch
+    intermediate_bar.close()
+
+    # Phase 3: time calibration batches from the middle
+    calibration_times: List[float] = []
+    cal_bar = tqdm(range(n_calibration), desc='Calibrating ETA', leave=False)
+    for j in cal_bar:
+        t0 = time.perf_counter()
+        batch = next(dl_iter)
+        yield mid_start + j, batch
+        calibration_times.append(time.perf_counter() - t0)
+    cal_bar.close()
+
+    avg_time = sum(calibration_times) / len(calibration_times)
+    remaining_start = mid_start + n_calibration
+    remaining_count = total - remaining_start
+    estimated_total_seconds = avg_time * remaining_count
+
+    # Phase 4: remaining batches with calibrated ETA
+    main_bar = tqdm(
+        range(remaining_count),
+        desc='Embedding batches',
+        bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]',
+    )
+    main_bar.set_postfix_str(f'ETA ~{estimated_total_seconds:.0f}s (calibrated)')
+    for k in main_bar:
+        batch = next(dl_iter)
+        yield remaining_start + k, batch
+    main_bar.close()
+
+
+class _SQLWriter:
+    """Context manager for async SQL embedding writes. Matches core/embed/storage.SQLEmbeddingWriter."""
+
+    def __init__(self, conn: sqlite3.Connection, queue_maxsize: int = 4) -> None:
+        self._conn = conn
+        self._queue: queue.Queue = queue.Queue(maxsize=queue_maxsize)
+        self._thread: Optional[threading.Thread] = None
+
+    def __enter__(self) -> "_SQLWriter":
+        self._thread = threading.Thread(target=self._writer_loop, daemon=True)
+        self._thread.start()
+        return self
+
+    def write_batch(self, rows: List[Tuple[str, bytes]]) -> None:
+        self._queue.put(rows)
+
+    def _writer_loop(self) -> None:
+        cursor = self._conn.cursor()
+        while True:
+            item = self._queue.get()
+            if item is None:
+                break
+            cursor.executemany("INSERT OR REPLACE INTO embeddings VALUES (?, ?)", item)
+            if self._queue.qsize() == 0:
+                self._conn.commit()
+        self._conn.commit()
+
+    def __exit__(self, *exc) -> None:
+        if self._thread is not None:
+            self._queue.put(None)
+            self._thread.join()
+            self._thread = None
+
+
+class Pooler:
+    def __init__(self, pooling_types: List[str]) -> None:
+        self.pooling_types = pooling_types
+        self.pooling_options: Dict[str, Callable] = {
+            'mean': self.mean_pooling,
+            'max': self.max_pooling,
+            'norm': self.norm_pooling,
+            'median': self.median_pooling,
+            'std': self.std_pooling,
+            'var': self.var_pooling,
+            'cls': self.cls_pooling,
+            'parti': self._pool_parti,
+        }
+
+    def _create_pooled_matrices_across_layers(self, attentions: torch.Tensor) -> torch.Tensor:
+        assert isinstance(attentions, torch.Tensor)
+        maxed_attentions = torch.max(attentions, dim=1)[0]
+        return maxed_attentions
+
+    def _page_rank(self, attention_matrix: np.ndarray, personalization: Optional[dict] = None, nstart: Optional[dict] = None, prune_type: str = "top_k_outdegree") -> Dict[int, float]:
+        G = self._convert_to_graph(attention_matrix)
+        if G.number_of_nodes() != attention_matrix.shape[0]:
+            raise Exception(
+                f"The number of nodes in the graph should be equal to the number of tokens in sequence! You have {G.number_of_nodes()} nodes for {attention_matrix.shape[0]} tokens.")
+        if G.number_of_edges() == 0:
+            raise Exception(f"You don't seem to have any attention edges left in the graph.")
+
+        return nx.pagerank(G, alpha=0.85, tol=1e-06, weight='weight', personalization=personalization, nstart=nstart, max_iter=100)
+
+    def _convert_to_graph(self, matrix: np.ndarray) -> nx.DiGraph:
+        G = nx.from_numpy_array(matrix, create_using=nx.DiGraph)
+        return G
+
+    def _calculate_importance_weights(self, dict_importance: Dict[int, float], attention_mask: Optional[torch.Tensor] = None) -> np.ndarray:
+        if attention_mask is not None:
+            for k in list(dict_importance.keys()):
+                if attention_mask[k] == 0:
+                    del dict_importance[k]
+
+        total = sum(dict_importance.values())
+        return np.array([v / total for _, v in dict_importance.items()])
+
+    def _pool_parti(self, emb: torch.Tensor, attentions: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        maxed_attentions = self._create_pooled_matrices_across_layers(attentions).numpy()
+        emb_pooled = []
+        for e, a, mask in zip(emb, maxed_attentions, attention_mask):
+            dict_importance = self._page_rank(a)
+            importance_weights = self._calculate_importance_weights(dict_importance, mask)
+            num_tokens = int(mask.sum().item())
+            emb_pooled.append(np.average(e[:num_tokens], weights=importance_weights, axis=0))
+        pooled = torch.tensor(np.array(emb_pooled))
+        return pooled
+
+    def mean_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        if attention_mask is None:
+            return emb.mean(dim=1)
+        else:
+            attention_mask = attention_mask.unsqueeze(-1)
+            return (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
+
+    def max_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        if attention_mask is None:
+            return emb.max(dim=1).values
+        else:
+            mask = attention_mask.unsqueeze(-1).bool()
+            return emb.masked_fill(~mask, float('-inf')).max(dim=1).values
+
+    def norm_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        if attention_mask is None:
+            return emb.norm(dim=1, p=2)
+        else:
+            attention_mask = attention_mask.unsqueeze(-1)
+            return (emb * attention_mask).norm(dim=1, p=2)
+
+    def median_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        if attention_mask is None:
+            return emb.median(dim=1).values
+        else:
+            mask = attention_mask.unsqueeze(-1).bool()
+            return emb.masked_fill(~mask, float('nan')).nanmedian(dim=1).values
+
+    def std_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        if attention_mask is None:
+            return emb.std(dim=1)
+        else:
+            var = self.var_pooling(emb, attention_mask, **kwargs)
+            return torch.sqrt(var)
+
+    def var_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        if attention_mask is None:
+            return emb.var(dim=1)
+        else:
+            attention_mask = attention_mask.unsqueeze(-1)
+            mean = (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
+            mean = mean.unsqueeze(1)
+            squared_diff = (emb - mean) ** 2
+            var = (squared_diff * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
+            return var
+
+    def cls_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+        return emb[:, 0, :]
+
+    def __call__(
+            self,
+            emb: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            attentions: Optional[torch.Tensor] = None
+        ) -> torch.Tensor:
+        if attention_mask is not None:
+            assert attention_mask.sum(dim=-1).min() > 0, (
+                "Pooler received samples with all-zero attention masks. "
+                "This causes NaN from division by zero. Filter empty inputs before pooling."
+            )
+        final_emb: List[torch.Tensor] = []
+        for pooling_type in self.pooling_types:
+            final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions))
+        return torch.cat(final_emb, dim=-1)
+
+
+class ProteinDataset(TorchDataset):
+    """Simple dataset for protein sequences."""
+    def __init__(self, sequences: List[str]) -> None:
+        self.sequences = sequences
+
+    def __len__(self) -> int:
+        return len(self.sequences)
+
+    def __getitem__(self, idx: int) -> str:
+        return self.sequences[idx]
+
+
+def parse_fasta(fasta_path: str) -> List[str]:
+    assert os.path.exists(fasta_path), f"FASTA file does not exist: {fasta_path}"
+    sequences = []
+    current_seq = []
+    with open(fasta_path, 'r') as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            if line.startswith('>'):
+                if current_seq:
+                    sequences.append(''.join(current_seq))
+                    current_seq = []
+            else:
+                current_seq.append(line)
+    if current_seq:
+        sequences.append(''.join(current_seq))
+    return sequences
+
+
+class EmbeddingMixin:
+    def _embed(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        raise NotImplementedError
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device of the model."""
+        return next(self.parameters()).device
+
+    def _read_sequences_from_db(self, db_path: str) -> Set[str]:
+        """Read sequences from SQLite database."""
+        with sqlite3.connect(db_path, timeout=30) as conn:
+            c = conn.cursor()
+            c.execute("SELECT sequence FROM embeddings")
+            return {row[0] for row in c.fetchall()}
+
+    def _ensure_embeddings_table(self, conn: sqlite3.Connection) -> None:
+        cursor = conn.cursor()
+        cursor.execute(
+            "CREATE TABLE IF NOT EXISTS embeddings ("
+            "sequence TEXT PRIMARY KEY, "
+            "embedding BLOB NOT NULL"
+            ")"
+        )
+        conn.commit()
+
+    def load_embeddings_from_pth(self, save_path: str) -> Dict[str, torch.Tensor]:
+        assert os.path.exists(save_path), f"Embedding file does not exist: {save_path}"
+        payload = torch.load(save_path, map_location="cpu", weights_only=True)
+        assert isinstance(payload, dict), "Expected .pth embeddings file to contain a dictionary."
+        for sequence, tensor in payload.items():
+            assert isinstance(sequence, str), "Expected embedding dictionary keys to be sequences (str)."
+            assert isinstance(tensor, torch.Tensor), "Expected embedding dictionary values to be tensors."
+        return payload
+
+    def load_embeddings_from_db(self, db_path: str, sequences: Optional[List[str]] = None) -> Dict[str, torch.Tensor]:
+        assert os.path.exists(db_path), f"Embedding database does not exist: {db_path}"
+        loaded: Dict[str, torch.Tensor] = {}
+        with sqlite3.connect(db_path, timeout=30) as conn:
+            self._ensure_embeddings_table(conn)
+            cursor = conn.cursor()
+            if sequences is None:
+                cursor.execute("SELECT sequence, embedding FROM embeddings")
+            else:
+                if len(sequences) == 0:
+                    return loaded
+                placeholders = ",".join(["?"] * len(sequences))
+                cursor.execute(
+                    f"SELECT sequence, embedding FROM embeddings WHERE sequence IN ({placeholders})",
+                    tuple(sequences),
+                )
+
+            rows = cursor.fetchall()
+            for row in rows:
+                sequence = row[0]
+                embedding_bytes = row[1]
+                loaded[sequence] = embedding_blob_to_tensor(embedding_bytes)
+        return loaded
+
+    def embed_dataset(
+        self,
+        sequences: Optional[List[str]] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+        batch_size: int = 2,
+        max_len: int = 512,
+        truncate: bool = True,
+        full_embeddings: bool = False,
+        embed_dtype: torch.dtype = torch.float32,
+        pooling_types: List[str] = ['mean'],
+        num_workers: int = 0,
+        sql: bool = False,
+        save: bool = True,
+        sql_db_path: str = 'embeddings.db',
+        save_path: str = 'embeddings.pth',
+        fasta_path: Optional[str] = None,
+        padding: str = 'max_length',
+        **kwargs,
+    ) -> Optional[Dict[str, torch.Tensor]]:
+        """
+        Embed a dataset of protein sequences.
+
+        Supports two modes:
+        - Tokenizer mode (ESM2/ESM++): provide `tokenizer` or use `self.tokenizer`.
+        - Sequence mode (E1): pass `tokenizer=None`, `_embed(sequences, return_attention_mask=True, **kwargs)` is used.
+
+        Sequences can be supplied as a list via `sequences`, parsed from a FASTA file via
+        `fasta_path`, or both (the two sources are combined). At least one must be provided.
+        """
+        if fasta_path is not None:
+            fasta_sequences = parse_fasta(fasta_path)
+            sequences = list(sequences or []) + fasta_sequences
+        assert sequences is not None and len(sequences) > 0, \
+            "Must provide at least one sequence via `sequences` or `fasta_path`."
+        sequences = list(set([seq[:max_len] if truncate else seq for seq in sequences]))
+        sequences = sorted(sequences, key=len, reverse=True)
+        hidden_size = self.config.hidden_size
+        pooler = Pooler(pooling_types) if not full_embeddings else None
+        if tokenizer is None and self.config.model_type != "E1":
+            tokenizer = self.tokenizer
+        tokenizer_mode = tokenizer is not None
+
+        # Resolve padding and compilation
+        dynamic = padding == 'longest'
+        compiled_model = maybe_compile(self, dynamic=dynamic)
+
+        if tokenizer_mode:
+            collate_fn = build_collator(tokenizer, padding=padding, max_length=max_len)
+            device = self.device
+        else:
+            collate_fn = None
+            device = None
+
+        def get_embeddings(residue_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+            assert isinstance(residue_embeddings, torch.Tensor)
+            if full_embeddings or residue_embeddings.ndim == 2:
+                return residue_embeddings
+            return pooler(residue_embeddings, attention_mask)
+
+        def iter_batches(to_embed: List[str]):
+            if tokenizer_mode:
+                assert collate_fn is not None
+                assert device is not None
+                dataset = ProteinDataset(to_embed)
+                dataloader = DataLoader(
+                    dataset,
+                    batch_size=batch_size,
+                    num_workers=num_workers,
+                    prefetch_factor=2 if num_workers > 0 else None,
+                    collate_fn=collate_fn,
+                    shuffle=False,
+                    pin_memory=True,
+                )
+                for i, batch in _make_embedding_progress(dataloader, padding):
+                    seqs = to_embed[i * batch_size:(i + 1) * batch_size]
+                    input_ids = batch['input_ids'].to(device)
+                    attention_mask = batch['attention_mask'].to(device)
+                    residue_embeddings = compiled_model._embed(input_ids, attention_mask)
+                    yield seqs, residue_embeddings, attention_mask
+            else:
+                for batch_start in tqdm(range(0, len(to_embed), batch_size), desc='Embedding batches'):
+                    seqs = to_embed[batch_start:batch_start + batch_size]
+                    batch_output = compiled_model._embed(seqs, return_attention_mask=True, **kwargs)
+                    assert isinstance(batch_output, tuple), "Sequence mode _embed must return (last_hidden_state, attention_mask)."
+                    assert len(batch_output) == 2, "Sequence mode _embed must return exactly two values."
+                    residue_embeddings, attention_mask = batch_output
+                    assert isinstance(attention_mask, torch.Tensor), "Sequence mode _embed must return attention_mask as a torch.Tensor."
+                    yield seqs, residue_embeddings, attention_mask
+
+        if sql:
+            # Step 1: DEDUPLICATE - check existing embeddings in SQL
+            conn = sqlite3.connect(sql_db_path, timeout=30, check_same_thread=False)
+            conn.execute('PRAGMA journal_mode=WAL')
+            conn.execute('PRAGMA busy_timeout=30000')
+            conn.execute('PRAGMA synchronous=OFF')
+            conn.execute('PRAGMA cache_size=-64000')
+            self._ensure_embeddings_table(conn)
+            already_embedded = self._read_sequences_from_db(sql_db_path)
+            to_embed = [seq for seq in sequences if seq not in already_embedded]
+            print(f"Found {len(already_embedded)} already embedded sequences in {sql_db_path}")
+            print(f"Embedding {len(to_embed)} new sequences")
+            if len(to_embed) > 0:
+                # Steps 4-7: BATCH+EMBED -> POOL/TRIM -> SERIALIZE -> WRITE (async)
+                with _SQLWriter(conn) as writer:
+                    with torch.inference_mode():
+                        for seqs, residue_embeddings, attention_mask in iter_batches(to_embed):
+                            embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
+                            if full_embeddings:
+                                batch_rows = []
+                                for seq, emb, mask in zip(seqs, embeddings, attention_mask):
+                                    batch_rows.append((seq, tensor_to_embedding_blob(emb[mask.bool()].reshape(-1, hidden_size))))
+                            else:
+                                blobs = batch_tensor_to_blobs(embeddings)
+                                batch_rows = list(zip(seqs, blobs))
+                            writer.write_batch(batch_rows)
+            conn.close()
+            return None
+
+        embeddings_dict = {}
+        if os.path.exists(save_path):
+            embeddings_dict = self.load_embeddings_from_pth(save_path)
+            to_embed = [seq for seq in sequences if seq not in embeddings_dict]
+            print(f"Found {len(embeddings_dict)} already embedded sequences in {save_path}")
+            print(f"Embedding {len(to_embed)} new sequences")
+        else:
+            to_embed = sequences
+            print(f"Embedding {len(to_embed)} new sequences")
+
+        if len(to_embed) > 0:
+            with torch.inference_mode():
+                for seqs, residue_embeddings, attention_mask in iter_batches(to_embed):
+                    embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
+                    for seq, emb, mask in zip(seqs, embeddings, attention_mask):
+                        if full_embeddings:
+                            emb = emb[mask.bool()].reshape(-1, hidden_size)
+                        embeddings_dict[seq] = emb.cpu()
+
+        if save:
+            torch.save(embeddings_dict, save_path)
+
+        return embeddings_dict
+
+
+if __name__ == "__main__":
+    # py -m pooler
+    pooler = Pooler(pooling_types=['max', 'parti'])
+    batch_size = 8
+    seq_len = 64
+    hidden_size = 128
+    num_layers = 12
+    emb = torch.randn(batch_size, seq_len, hidden_size)
+    attentions = torch.randn(batch_size, num_layers, seq_len, seq_len)
+    attention_mask = torch.ones(batch_size, seq_len)
+    y = pooler(emb=emb, attention_mask=attention_mask, attentions=attentions)
+    print(y.shape)
+
+"""Shared attention infrastructure for all FastPLMs models.
+
+Contains: AttentionBackend enum, backend resolution, mask creation,
+flex attention helpers, flash kernel detection/dispatch, and pad/unpad utilities.
+"""
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from einops import rearrange
+
+try:
+    from torch.nn.attention.flex_attention import create_block_mask, flex_attention, BlockMask
+except ImportError:
+    create_block_mask = None
+    flex_attention = None
+    BlockMask = None
+
+_compiled_flex_attention = None
+
+
+def _get_flex_attention_fn():
+    """Return flex_attention callable: compiled (fused kernel) by default, or eager when debug flag is set."""
+    global _compiled_flex_attention
+    if flex_attention is None:
+        return None
+    flex_mod = torch.nn.attention.flex_attention
+    if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False):
+        return flex_attention
+    if _compiled_flex_attention is None:
+        _compiled_flex_attention = torch.compile(
+            flex_attention,
+            dynamic=False,
+        )
+    return _compiled_flex_attention
+
+
+### Kernels Flash Attention Detection
+def _infer_kernels_flash_variant(kernel) -> Optional[str]:
+    if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"):
+        return "flash_attn2"
+    if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"):
+        return "flash_attn3"
+    return None
+
+
+def _try_get_kernels_flash():
+    try:
+        from kernels import get_kernel
+    except ImportError:
+        return None, None
+
+    flash_kernel = None
+    flash_kernel_variant = None
+    try:
+        flash_kernel = get_kernel("kernels-community/flash-attn3")
+        flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
+        assert flash_kernel_variant is not None, "Loaded flash-attn3 kernel does not expose a supported API."
+    except Exception:
+        try:
+            flash_kernel = get_kernel("kernels-community/flash-attn2")
+            flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
+            assert flash_kernel_variant is not None, "Loaded flash-attn2 kernel does not expose a supported API."
+        except Exception:
+            flash_kernel = None
+            flash_kernel_variant = None
+    return flash_kernel, flash_kernel_variant
+
+
+_FLASH_KERNELS_LOADED = False
+FLASH_KERNEL = None
+FLASH_KERNEL_VARIANT = None
+
+
+def _ensure_flash_kernels_loaded():
+    global _FLASH_KERNELS_LOADED, FLASH_KERNEL, FLASH_KERNEL_VARIANT
+    if _FLASH_KERNELS_LOADED:
+        return
+    _FLASH_KERNELS_LOADED = True
+    FLASH_KERNEL, FLASH_KERNEL_VARIANT = _try_get_kernels_flash()
+
+
+def _kernels_flash_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    causal: bool = False,
+    softmax_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Flash-attention forward, optionally overriding the softmax scale.
+
+    When `softmax_scale is None`, the flash kernel applies its default
+    `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already
+    pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold).
+    Failing to override when Q is pre-scaled produces DOUBLE scaling and
+    catastrophic downstream drift -- on DPLM-150M (30 layers) this was observed
+    as pooled-embedding cosine ~-0.12 and argmax agreement ~0.27 vs sdpa.
+    """
+    assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
+    if FLASH_KERNEL_VARIANT == "flash_attn2":
+        return FLASH_KERNEL.fwd(
+            q=query_states, k=key_states, v=value_states,
+            softmax_scale=softmax_scale, is_causal=causal,
+        )[0]
+    if FLASH_KERNEL_VARIANT == "flash_attn3":
+        try:
+            output = FLASH_KERNEL.flash_attn_func(
+                q=query_states, k=key_states, v=value_states,
+                softmax_scale=softmax_scale, causal=causal,
+            )
+        except TypeError:
+            output = FLASH_KERNEL.flash_attn_func(
+                query_states, key_states, value_states,
+                0.0, softmax_scale, causal,
+            )
+        if isinstance(output, tuple):
+            return output[0]
+        return output
+    raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}")
+
+
+def _kernels_flash_varlen_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_in_batch_q: int,
+    max_seqlen_in_batch_k: int,
+    causal: bool = False,
+    softmax_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Varlen flash-attention forward, optionally overriding the softmax scale.
+
+    See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be
+    passed when Q has been pre-scaled by the caller.
+    """
+    assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
+    if FLASH_KERNEL_VARIANT == "flash_attn2":
+        return FLASH_KERNEL.varlen_fwd(
+            q=query_states, k=key_states, v=value_states,
+            cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k,
+            softmax_scale=softmax_scale, is_causal=causal,
+        )[0]
+    if FLASH_KERNEL_VARIANT == "flash_attn3":
+        try:
+            output = FLASH_KERNEL.flash_attn_varlen_func(
+                q=query_states, k=key_states, v=value_states,
+                cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q, max_seqlen_k=max_seqlen_in_batch_k,
+                softmax_scale=softmax_scale, causal=causal,
+            )
+        except TypeError:
+            output = FLASH_KERNEL.flash_attn_varlen_func(
+                query_states, key_states, value_states,
+                cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_in_batch_q, max_seqlen_in_batch_k,
+                0.0, softmax_scale, causal,
+            )
+        if isinstance(output, tuple):
+            return output[0]
+        return output
+    raise AssertionError(f"Unsupported kernels flash attention variant: {FLASH_KERNEL_VARIANT}")
+
+
+### Unpad / Pad helpers for varlen flash attention
+class IndexFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, indices) -> torch.Tensor:
+        ctx.save_for_backward(indices)
+        assert input.ndim >= 2
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        return torch.gather(
+            rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim)
+        ).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, grad_output) -> Tuple[torch.Tensor, None]:
+        (indices,) = ctx.saved_tensors
+        assert grad_output.ndim >= 2
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, "b ... -> b (...)")
+        grad_input = torch.zeros(
+            [ctx.first_axis_dim, grad_output.shape[1]], device=grad_output.device, dtype=grad_output.dtype
+        )
+        grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor:
+        ctx.save_for_backward(indices)
+        assert indices.ndim == 1
+        assert values.ndim >= 2
+        output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)
+        output[indices] = values
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output) -> Tuple[torch.Tensor, None, None]:
+        (indices,) = ctx.saved_tensors
+        return grad_output[indices], None, None
+
+
+index_first_axis = IndexFirstAxis.apply
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+def pad_input(hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+def _unpad_input(
+    query_layer: torch.Tensor,
+    key_layer: torch.Tensor,
+    value_layer: torch.Tensor,
+    attention_mask_2d: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor], Tuple[int, int]]:
+    batch_size, seq_len, num_heads, head_dim = query_layer.shape
+    seqlens = attention_mask_2d.sum(dim=1).int()
+    cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))
+    max_seqlen = int(seqlens.max().item())
+    indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()
+    query_layer = index_first_axis(query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
+    key_layer = index_first_axis(key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
+    value_layer = index_first_axis(value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices)
+    return query_layer, key_layer, value_layer, indices, (cu_seqlens, cu_seqlens), (max_seqlen, max_seqlen)
+
+
+def kernels_flash_attention_func(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attention_mask_2d: Optional[torch.Tensor] = None,
+    causal: bool = False,
+    softmax_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Public flash-attention entry point with optional padding handling.
+
+    `softmax_scale`:
+        None -> kernel applies its default `1 / sqrt(head_dim)`.
+        float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled
+        by the caller).
+
+    IMPORTANT: if your family multiplies Q by `1/sqrt(head_dim)` before calling
+    this function (as ESM2, DPLM, DPLM2, E1, and ESMFold do) you MUST pass
+    `softmax_scale=1.0`. Otherwise the kernel applies its default scale ON TOP
+    of the caller's, producing effective scale `1/head_dim` and catastrophic
+    downstream drift that compounds across layers.
+    """
+    assert FLASH_KERNEL is not None, "Kernel Flash Attention is not available in this environment."
+    if not causal and attention_mask_2d is not None:
+        batch_size, q_len = query_states.shape[:2]
+        (
+            query_states, key_states, value_states,
+            indices_q, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k),
+        ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d)
+        attn_output_unpad = _kernels_flash_varlen_forward(
+            query_states=query_states, key_states=key_states, value_states=value_states,
+            cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_in_batch_q=max_seqlen_q, max_seqlen_in_batch_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+        )
+        return pad_input(attn_output_unpad, indices_q, batch_size, q_len)
+    else:
+        return _kernels_flash_forward(
+            query_states=query_states, key_states=key_states, value_states=value_states,
+            causal=causal, softmax_scale=softmax_scale,
+        )
+
+
+### Attention Backend Enum & Resolution
+class AttentionBackend(Enum):
+    AUTO = "auto"
+    KERNELS_FLASH = "kernels_flash"
+    FLEX = "flex"
+    SDPA = "sdpa"
+
+
+VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend)
+
+
+_BACKEND_CONFIRMED = False
+
+
+def resolve_attention_backend(requested_backend: str) -> AttentionBackend:
+    global _BACKEND_CONFIRMED
+    assert requested_backend in VALID_ATTENTION_BACKENDS, (
+        f"Unsupported attention backend: {requested_backend}. Expected one of {VALID_ATTENTION_BACKENDS}."
+    )
+    if requested_backend in (AttentionBackend.AUTO.value, AttentionBackend.KERNELS_FLASH.value):
+        _ensure_flash_kernels_loaded()
+    if requested_backend == AttentionBackend.AUTO.value:
+        if FLASH_KERNEL is not None:
+            resolved = AttentionBackend.KERNELS_FLASH
+        elif flex_attention is not None:
+            resolved = AttentionBackend.FLEX
+        else:
+            resolved = AttentionBackend.SDPA
+    elif requested_backend == AttentionBackend.KERNELS_FLASH.value:
+        assert FLASH_KERNEL is not None, "Kernels Flash Attention is not available in this environment."
+        resolved = AttentionBackend.KERNELS_FLASH
+    elif requested_backend == AttentionBackend.FLEX.value:
+        assert flex_attention is not None, "Flex Attention is not available in this environment."
+        resolved = AttentionBackend.FLEX
+    elif requested_backend == AttentionBackend.SDPA.value:
+        resolved = AttentionBackend.SDPA
+    else:
+        raise AssertionError(f"Unsupported attention backend: {requested_backend}")
+    if not _BACKEND_CONFIRMED:
+        print(f"Attention backend: config='{requested_backend}' -> resolved='{resolved.value}'")
+        _BACKEND_CONFIRMED = True
+    return resolved
+
+
+@torch.compiler.disable
+def get_attention_mask(
+    effective_backend: AttentionBackend,
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[BlockMask]]:
+    """Build padding masks once for all encoder layers.
+
+    Returns (attention_mask_2d, attention_mask_4d, flex_block_mask).
+    """
+    if attention_mask is None:
+        return None, None, None
+
+    attention_mask_2d = attention_mask.bool()
+
+    if effective_backend == AttentionBackend.KERNELS_FLASH:
+        return attention_mask_2d, None, None
+
+    if effective_backend == AttentionBackend.FLEX:
+        assert create_block_mask is not None, "Flex attention backend requested but torch.create_block_mask is unavailable."
+        valid_lens = attention_mask_2d.sum(dim=-1)
+
+        def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+            return (q_idx < valid_lens[batch_idx]) & (kv_idx < valid_lens[batch_idx])
+
+        flex_block_mask = create_block_mask(mask_mod, batch_size, 1, seq_len, seq_len, device=device)
+        return attention_mask_2d, None, flex_block_mask
+
+    # SDPA / manual -- only mask the key dimension so padding query positions attend to
+    # real keys and produce valid (non-NaN) outputs instead of NaN from softmax(-inf,...,-inf).
+    attention_mask_4d = attention_mask_2d[:, None, None, :]
+    return attention_mask_2d, attention_mask_4d, None
+
+
+def bool_to_additive_mask(
+    bool_mask: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid).
+
+    Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))`
+    directly on a bool tensor returns a bool tensor -- because `-inf` casts to `True` -- and
+    silently drops the mask entirely. Always allocate a float tensor first, then fill it.
+    This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask.
+    """
+    assert bool_mask.dtype == torch.bool, (
+        f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}"
+    )
+    additive = torch.zeros_like(bool_mask, dtype=dtype)
+    additive.masked_fill_(bool_mask.logical_not(), float("-inf"))
+    return additive
+
 """
 ESM++ model implementation.
 
@@ -28,20 +1039,6 @@ from tokenizers.processors import TemplateProcessing
 from transformers import PreTrainedModel, PreTrainedTokenizerFast, PretrainedConfig
 from transformers.modeling_outputs import ModelOutput
 
-try:
-    from fastplms.attention import (
-        AttentionBackend, VALID_ATTENTION_BACKENDS,
-        resolve_attention_backend, get_attention_mask,
-        _get_flex_attention_fn,
-        _ensure_flash_kernels_loaded, FLASH_KERNEL, FLASH_KERNEL_VARIANT,
-        _kernels_flash_forward, _kernels_flash_varlen_forward,
-        kernels_flash_attention_func,
-        index_first_axis, index_put_first_axis, pad_input, _unpad_input,
-        create_block_mask, flex_attention, BlockMask,
-    )
-    from fastplms.embedding_mixin import Pooler, EmbeddingMixin, ProteinDataset, parse_fasta, build_collator
-except ImportError:
-    pass  # Running as HF Hub composite; shared definitions are above
 
 
 class ESMplusplusConfig(PretrainedConfig):