Upload modeling_ankh.py with huggingface_hub

modeling_ankh.py

@@ -1,751 +1,5 @@
 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 os
-import sqlite3
-import networkx as nx
-import numpy as np
-import torch
-from tqdm.auto import tqdm
-from typing import Callable, Dict, List, Optional, Set
-from torch.utils.data import DataLoader
-from torch.utils.data import Dataset as TorchDataset
-from transformers import PreTrainedTokenizerBase
-
-
-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]:
-        # Run PageRank on the attention matrix converted to a graph.
-        # Raises exceptions if the graph doesn't match the token sequence or has no edges.
-        # Returns the PageRank scores for each token node.
-        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:
-        # Convert a matrix (e.g., attention scores) to a directed graph using networkx.
-        # Each element in the matrix represents a directed edge with a weight.
-        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:
-        # Remove keys where attention_mask is 0
-        if attention_mask is not None:
-            for k in list(dict_importance.keys()):
-                if attention_mask[k] == 0:
-                    del dict_importance[k]
-
-        #dict_importance[0] # remove cls
-        #dict_importance[-1] # remove eos
-        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: # (b, L, d) -> (b, d)
-        maxed_attentions = self._create_pooled_matrices_across_layers(attentions).numpy()
-        # emb is (b, L, d), maxed_attentions is (b, L, L)
-        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: # (b, L, d) -> (b, d)
-        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: # (b, L, d) -> (b, d)
-        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: # (b, L, d) -> (b, d)
-        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: # (b, L, d) -> (b, d)
-        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: # (b, L, d) -> (b, d)
-        if attention_mask is None:
-            return emb.std(dim=1)
-        else:
-            # Compute variance correctly over non-masked positions, then take sqrt
-            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: # (b, L, d) -> (b, d)
-        if attention_mask is None:
-            return emb.var(dim=1)
-        else:
-            # Correctly compute variance over only non-masked positions
-            attention_mask = attention_mask.unsqueeze(-1)  # (b, L, 1)
-            # Compute mean over non-masked positions
-            mean = (emb * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)  # (b, d)
-            mean = mean.unsqueeze(1)  # (b, 1, d)
-            # Compute squared differences from mean, only over non-masked positions
-            squared_diff = (emb - mean) ** 2  # (b, L, d)
-            # Sum squared differences over non-masked positions and divide by count
-            var = (squared_diff * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)  # (b, d)
-            return var
-
-    def cls_pooling(self, emb: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor: # (b, L, d) -> (b, d)
-        return emb[:, 0, :]
-
-    def __call__(
-            self,
-            emb: torch.Tensor,
-            attention_mask: Optional[torch.Tensor] = None,
-            attentions: Optional[torch.Tensor] = None
-        ) -> torch.Tensor: # [mean, max]
-        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)) # (b, d)
-        return torch.cat(final_emb, dim=-1) # (b, n_pooling_types * d)
-
-
-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 build_collator(tokenizer: PreTrainedTokenizerBase) -> Callable[[List[str]], Dict[str, torch.Tensor]]:
-    def _collate_fn(sequences: List[str]) -> Dict[str, torch.Tensor]:
-        return tokenizer(sequences, return_tensors="pt", padding='longest')
-    return _collate_fn
-
-
-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."""
-        sequences = []
-        with sqlite3.connect(db_path) as conn:
-            c = conn.cursor()
-            c.execute("SELECT sequence FROM embeddings")
-            while True:
-                row = c.fetchone()
-                if row is None:
-                    break
-                sequences.append(row[0])
-        return set(sequences)
-
-    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, "
-            "shape TEXT, "
-            "dtype TEXT"
-            ")"
-        )
-        cursor.execute("PRAGMA table_info(embeddings)")
-        rows = cursor.fetchall()
-        column_names = [row[1] for row in rows]
-        if "shape" not in column_names:
-            cursor.execute("ALTER TABLE embeddings ADD COLUMN shape TEXT")
-        if "dtype" not in column_names:
-            cursor.execute("ALTER TABLE embeddings ADD COLUMN dtype TEXT")
-        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) as conn:
-            self._ensure_embeddings_table(conn)
-            cursor = conn.cursor()
-            if sequences is None:
-                cursor.execute("SELECT sequence, embedding, shape, dtype FROM embeddings")
-            else:
-                if len(sequences) == 0:
-                    return loaded
-                placeholders = ",".join(["?"] * len(sequences))
-                cursor.execute(
-                    f"SELECT sequence, embedding, shape, dtype FROM embeddings WHERE sequence IN ({placeholders})",
-                    tuple(sequences),
-                )
-
-            rows = cursor.fetchall()
-            for row in rows:
-                sequence = row[0]
-                embedding_bytes = row[1]
-                shape_text = row[2]
-                dtype_text = row[3]
-                assert shape_text is not None, "Missing shape metadata in embeddings table."
-                assert dtype_text is not None, "Missing dtype metadata in embeddings table."
-                shape_values = [int(value) for value in shape_text.split(",") if len(value) > 0]
-                assert len(shape_values) > 0, f"Invalid shape metadata for sequence: {sequence}"
-                expected_size = int(np.prod(shape_values))
-                np_dtype = np.dtype(dtype_text)
-                array = np.frombuffer(embedding_bytes, dtype=np_dtype)
-                assert array.size == expected_size, f"Shape mismatch while reading sequence: {sequence}"
-                reshaped = array.copy().reshape(tuple(shape_values))
-                loaded[sequence] = torch.from_numpy(reshaped)
-        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,
-        **kwargs,
-    ) -> Optional[Dict[str, torch.Tensor]]:
-        """
-        Embed a dataset of protein sequences.
-
-        Supports two modes:
-        - Tokenizer mode (ESM2/ESM++): provide `tokenizer`, `_embed(input_ids, attention_mask)` is used.
-        - 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
-        tokenizer_mode = tokenizer is not None
-        if tokenizer_mode:
-            collate_fn = build_collator(tokenizer)
-            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, collate_fn=collate_fn, shuffle=False)
-                for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc='Embedding batches'):
-                    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 = self._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 = self._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:
-            conn = sqlite3.connect(sql_db_path)
-            self._ensure_embeddings_table(conn)
-            c = conn.cursor()
-            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:
-                with torch.no_grad():
-                    for i, (seqs, residue_embeddings, attention_mask) in enumerate(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)
-                            emb_np = emb.cpu().numpy()
-                            emb_shape = ",".join([str(dim) for dim in emb_np.shape])
-                            emb_dtype = str(emb_np.dtype)
-                            c.execute(
-                                "INSERT OR REPLACE INTO embeddings (sequence, embedding, shape, dtype) VALUES (?, ?, ?, ?)",
-                                (seq, emb_np.tobytes(), emb_shape, emb_dtype),
-                            )
-                        if tokenizer_mode and (i + 1) % 100 == 0:
-                            conn.commit()
-                conn.commit()
-            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.no_grad():
-                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 functools import partial
-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.
-
-    Uses kernel_options={"BACKEND": "FLASH"} to prefer Flash Attention 4 (FA4)
-    on Hopper/Blackwell GPUs (PyTorch 2.11+). Automatically falls back to Triton
-    on older hardware.
-    """
-    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(
-            partial(flex_attention, kernel_options={"BACKEND": "FLASH"}),
-            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,
-) -> torch.Tensor:
-    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, 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, causal=causal)
-        except TypeError:
-            output = FLASH_KERNEL.flash_attn_func(query_states, key_states, value_states, 0.0, None, 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,
-) -> torch.Tensor:
-    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,
-            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,
-                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, None, 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,
-) -> torch.Tensor:
-    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,
-        )
-        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,
-        )
-
-
-### 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
-
 import math
 
 import torch
@@ -756,6 +10,16 @@ from dataclasses import dataclass
 from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer
 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,
+        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
 
 
 # ---------------------------------------------------------------------------