from __future__ import annotations import hashlib import os import sys import torch import torch.nn as nn import torch.nn.functional as F from collections import defaultdict from contextvars import ContextVar from dataclasses import dataclass from enum import Enum from typing import Any, ClassVar, TypedDict from tqdm.auto import tqdm from transformers import PretrainedConfig, PreTrainedModel from transformers.activations import ACT2FN from transformers.modeling_outputs import ModelOutput from transformers.utils import logging try: from fastplms.attention import ( AttentionBackend, BlockMask, FastPLMsAttentionMixin, resolve_attention_backend, resolve_attention_backend_for_call, ) from fastplms.embeddings import ( EmbeddingBatch, EmbeddingMixin, EmbeddingResult, Pooler, embed_dataset, select_hidden_state_embeddings, ) from fastplms.models.ttt import FastPLMTestTimeTrainingMixin except ModuleNotFoundError as error: _COMPOSITE_REQUIRED_NAMES = ( "AttentionBackend", "BlockMask", "EmbeddingBatch", "EmbeddingMixin", "EmbeddingResult", "FastPLMsAttentionMixin", "FastPLMTestTimeTrainingMixin", "Pooler", "embed_dataset", "resolve_attention_backend", "resolve_attention_backend_for_call", "select_hidden_state_embeddings", ) if error.name != "fastplms" or any( name not in globals() for name in _COMPOSITE_REQUIRED_NAMES ): raise # Legacy flat Hub composites define every shared symbol above this block. from .attention import ( # noqa: F401 _document_ids, _get_unpad_data, _unpad_input, block_mask_creator, block_min_max_seq_ids, build_block_causal_mask_4d, build_within_seq_mask_4d, create_block_causal_mask_optimized, create_within_seq_block_mask, direct_block_mask, doc_id_mask, flex_attention_func, get_overlapping_blocks, kernels_flash_attention_func, varlen_flex_attention_func, ) from .cache import DynamicCache, KVCache # noqa: F401 from .preparation import ( # noqa: F401 BOS_TOKEN_ID, E1_TOKENIZER_REPO_ID, E1_VOCAB_SIZE, EOS_TOKEN_ID, PAD_TOKEN_ID, DataPrepConfig, E1BatchPreparer, _load_tokenizer_file, get_context, get_tokenizer, ) from .retrieval import ( # noqa: F401 COLABFOLD_HOST, DEFAULT_EMBED_MAX_TOKENS, DEFAULT_EMBED_SIMILARITY, DEFAULT_MAX_CONTEXT_TOKENS, DEFAULT_SIMILARITY_THRESHOLDS, DOCKER_IMAGE, E1_MSA_SAMPLING_SOURCE_REVISION, LOWERCASE_CHARS, ColabFoldSearcher, ContextCache, ContextSpecification, E1Prediction, HomologueSearcher, IdSequence, IndexedSequence, _ColabFoldResponse, _E1ContextPredictor, _forward_for_embedding, _make_homologue_searcher, _pool_hidden_states, _safe_extract_tar, _sequence_output_dir, _strip_a3m_insertions, build_context_specifications, compute_ppll, convert_to_tensor, get_context_id, get_msa_for_sequence, get_num_neighbors, get_query_from_a3m, get_similarity_to_query, load_msa_dir, load_msa_from_hf, parse_msa, read_fasta_sequences, sample_context, sample_contexts_for_msa, sample_multiple_contexts, write_fasta_sequences, ) def _get_logger(): """Resolve the Transformers logger only when a runtime path emits a message.""" return logging.get_logger(__name__) _TOKENIZER_LOAD_CONTEXT: ContextVar[dict[str, Any] | None] = ContextVar( "fastplms_e1_tokenizer_load_context", default=None, ) class E1Config(PretrainedConfig): model_type = "E1" keys_to_ignore_at_inference: ClassVar[list[str]] = ["past_key_values"] def __init__( # type: ignore self, # Model architecture/initialization vocab_size=None, hidden_size=4096, intermediate_size=16384, gated_mlp=False, num_hidden_layers=40, num_attention_heads=32, num_key_value_heads=8, hidden_act="silu", rms_norm_eps=1e-5, initializer_range=0.02, dtype="bfloat16", gradient_checkpointing=False, no_ffn_gradient_checkpointing=False, use_cache=False, # Tokenization pad_token_id=None, bos_token_id=None, eos_token_id=None, tie_word_embeddings=False, # Attention implementation & rotary positional embeddings global_attention_every_n_layers=0, max_num_sequences=512, max_num_positions_within_seq=8192, max_num_positions_global=1024 * 128, rope_theta_within_seq=10000.0, rope_theta_global=100000.0, clip_qkv=None, attn_backend=None, **kwargs, ) -> None: super().__init__( pad_token_id=PAD_TOKEN_ID, bos_token_id=BOS_TOKEN_ID, eos_token_id=EOS_TOKEN_ID, tie_word_embeddings=tie_word_embeddings, dtype=dtype, **kwargs, ) self.hidden_size = hidden_size if intermediate_size is None: intermediate_size = 3 * hidden_size if gated_mlp else 4 * hidden_size self.intermediate_size = intermediate_size self.gated_mlp = gated_mlp self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.max_num_positions_within_seq = max_num_positions_within_seq self.max_num_positions_global = max_num_positions_global # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.rope_theta_within_seq = rope_theta_within_seq self.rope_theta_global = rope_theta_global self.max_num_sequences = max_num_sequences if clip_qkv is not None and clip_qkv <= 0: raise ValueError(f"clip_qkv must be positive when provided, got {clip_qkv}.") self.clip_qkv = clip_qkv self.global_attention_every_n_layers = global_attention_every_n_layers self.vocab_size = E1_VOCAB_SIZE self.gradient_checkpointing = gradient_checkpointing self.no_ffn_gradient_checkpointing = no_ffn_gradient_checkpointing if not isinstance(use_cache, bool): raise TypeError("use_cache must be a boolean.") self.use_cache = use_cache self.attn_backend = attn_backend if vocab_size is not None: if vocab_size < self.vocab_size: _get_logger().warning( f"Using vocab_size {vocab_size} smaller than {self.vocab_size} " "from the tokenizer contract." ) self.vocab_size = vocab_size elif vocab_size > self.vocab_size: _get_logger().warning( f"Using vocab_size {vocab_size} instead of smaller {self.vocab_size} " "from E1 tokenizer contract." ) self.vocab_size = vocab_size if pad_token_id is not None and pad_token_id != self.pad_token_id: _get_logger().warning( f"Ignoring pad_token_id. Using {self.pad_token_id} from E1 tokenizer contract" ) if bos_token_id is not None and bos_token_id != self.bos_token_id: _get_logger().warning( f"Ignoring bos_token_id. Using {self.bos_token_id} from E1 tokenizer contract" ) if eos_token_id is not None and eos_token_id != self.eos_token_id: _get_logger().warning( f"Ignoring eos_token_id. Using {self.eos_token_id} from E1 tokenizer contract" ) class AttentionLayerType(Enum): WITHIN_SEQ = "within_seq" GLOBAL = "global" class AttentionArgs(TypedDict, total=False): within_seq_block_mask: BlockMask | None block_causal_block_mask: BlockMask | None within_seq_mask_4d: torch.Tensor | None block_causal_mask_4d: torch.Tensor | None def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: """This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) """ batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_key_value_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) class RotaryPositionalEmbedding(nn.Module): def __init__( self, dim: int, max_position_embeddings: int = 2048, base: int = 10000, device: torch.device | None = None, ) -> None: super().__init__() self.dim = dim self.base = base self.max_position_embeddings = max_position_embeddings # Transformers may instantiate modules on the meta device while loading # a checkpoint. Precomputed non-persistent buffers would then be # materialized without values. Empty buffers make initialization lazy # and deterministic on the first real-device forward. empty = torch.empty(0, dtype=torch.float32, device=device) self.register_buffer("inv_freq", empty, persistent=False) self.register_buffer("cos_cached", empty.clone(), persistent=False) self.register_buffer("sin_cached", empty.clone(), persistent=False) self.max_seq_len_cached = 0 @staticmethod def rotate_half(x: torch.Tensor) -> torch.Tensor: """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def _set_sin_cos_cache(self, seq_len: int, device: torch.device) -> None: # Compute angles in FP32, matching the official cache constructed before # the model is converted to its inference dtype. self.max_seq_len_cached = seq_len inv_freq = self.base ** -( torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim ) # (d / 2,) self.inv_freq = inv_freq t = torch.arange(seq_len, device=device, dtype=torch.float32) # (l,) angles = torch.outer(t, inv_freq) # (l, d / 2) angles = torch.cat((angles, angles), dim=1) # (l, d) self.cos_cached = angles.cos() self.sin_cached = angles.sin() def forward( self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.LongTensor, seq_len: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: # q, k: (b, l, h, d) device, dtype = q.device, q.dtype seq_len = position_ids.max().item() + 1 if seq_len is None else seq_len if seq_len > self.max_seq_len_cached: self._set_sin_cos_cache(seq_len=seq_len, device=device) # Selecting by position inserts a head axis for broadcasting. idxs = position_ids.to(device) cos = self.cos_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs] # (b, l, 1, d) sin = self.sin_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs] # (b, l, 1, d) # Apply the real and imaginary parts of the rotary transform to Q and K. # Both halves reuse C and S, so rotate_half supplies the cross terms. q_embed = (q * cos) + (self.rotate_half(q) * sin) k_embed = (k * cos) + (self.rotate_half(k) * sin) return q_embed, k_embed class Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper.""" def __init__(self, config: E1Config, layer_idx: int) -> None: super().__init__() self.config = config self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_kv_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_kv_heads self.max_num_seqs = config.max_num_sequences self.clip_qkv = config.clip_qkv if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError( f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) if self.config.global_attention_every_n_layers > 0: self.layer_type = ( AttentionLayerType.GLOBAL if (self.layer_idx + 1) % self.config.global_attention_every_n_layers == 0 else AttentionLayerType.WITHIN_SEQ ) else: self.layer_type = AttentionLayerType.WITHIN_SEQ self.rope_theta = ( config.rope_theta_within_seq if self.layer_type == AttentionLayerType.WITHIN_SEQ else config.rope_theta_global ) self.max_position_embeddings = ( config.max_num_positions_within_seq if self.layer_type == AttentionLayerType.WITHIN_SEQ else config.max_num_positions_global ) self.rotary_emb = RotaryPositionalEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) self.attn_backend = resolve_attention_backend(config.attn_backend) def prepare_qkv( self, hidden_states: torch.Tensor, position_ids: torch.LongTensor, past_key_value: DynamicCache | None = None, use_cache: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # hidden_states: (b, l, d); position_ids: (b, l) bsz, q_len, _ = hidden_states.size() query_states: torch.Tensor = self.q_proj(hidden_states) key_states: torch.Tensor = self.k_proj(hidden_states) val_states: torch.Tensor = self.v_proj(hidden_states) query_states = query_states.view( bsz, q_len, self.num_heads, self.head_dim ) # (b, l, h, d_h) key_states = key_states.view( bsz, q_len, self.num_kv_heads, self.head_dim ) # (b, l, h_kv, d_h) val_states = val_states.view( bsz, q_len, self.num_kv_heads, self.head_dim ) # (b, l, h_kv, d_h) if self.clip_qkv is not None: query_states = query_states.clamp(-self.clip_qkv, self.clip_qkv) key_states = key_states.clamp(-self.clip_qkv, self.clip_qkv) val_states = val_states.clamp(-self.clip_qkv, self.clip_qkv) query_states, key_states = self.rotary_emb(query_states, key_states, position_ids) if use_cache and past_key_value is not None: key_states, val_states = past_key_value.update(key_states, val_states, self.layer_idx) input_dtype = query_states.dtype if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_dtype("cuda") else: target_dtype = self.q_proj.weight.dtype if input_dtype != target_dtype: _get_logger().warning_once( f"The input hidden states seems to be silently casted in {input_dtype}. " f"This might be because you have upcasted embedding or layer norm layers " f"in {input_dtype}. We will cast back the input in {target_dtype}." ) query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) val_states = val_states.to(target_dtype) return query_states, key_states, val_states def forward( self, hidden_states: torch.Tensor, within_seq_position_ids: torch.LongTensor, global_position_ids: torch.LongTensor, sequence_ids: torch.LongTensor, attention_args: AttentionArgs | None = None, past_key_value: DynamicCache | None = None, output_attentions: bool = False, output_s_max: bool = False, use_cache: bool = False, effective_backend: AttentionBackend | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]: # hidden_states: (b, l, d); position and sequence IDs: (b, l) is_cache_prefilled = ( use_cache and past_key_value is not None and past_key_value.get_seq_length(self.layer_idx) > 0 ) query_states, key_states, val_states = self.prepare_qkv( hidden_states=hidden_states, position_ids=within_seq_position_ids if self.layer_type == AttentionLayerType.WITHIN_SEQ else global_position_ids, past_key_value=past_key_value, use_cache=use_cache, ) attn_output, attn_weights, s_max = self._attn( query_states=query_states, key_states=key_states, val_states=val_states, sequence_ids=sequence_ids, attention_args=attention_args, output_attentions=output_attentions, output_s_max=output_s_max, is_cache_prefilled=is_cache_prefilled, effective_backend=effective_backend, ) attn_output = self.o_proj(attn_output) return attn_output, attn_weights, past_key_value, s_max def _attn( self, query_states: torch.Tensor, key_states: torch.Tensor, val_states: torch.Tensor, sequence_ids: torch.Tensor, attention_args: AttentionArgs | None = None, output_attentions: bool = False, output_s_max: bool = False, is_cache_prefilled: bool = False, effective_backend: AttentionBackend | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: # A filled cache changes the implementation shape, not the layer's # biological attention contract. Global layers must retain the cached # context, while within-sequence layers consume only the newly appended # sequence. This matches the pinned E1 inference implementation. effective_layer_type = self.layer_type if effective_backend is None: effective_backend = resolve_attention_backend_for_call( self.attn_backend, output_attentions=output_attentions, ) if output_attentions: return self._manual_attn( query_states, key_states, val_states, sequence_ids=sequence_ids, attention_args=attention_args, effective_layer_type=effective_layer_type, output_s_max=output_s_max, is_cache_prefilled=is_cache_prefilled, ) if effective_backend == AttentionBackend.EAGER: attn_output, _, s_max = self._manual_attn( query_states, key_states, val_states, sequence_ids=sequence_ids, attention_args=attention_args, effective_layer_type=effective_layer_type, output_s_max=output_s_max, is_cache_prefilled=is_cache_prefilled, ) return attn_output, None, s_max if effective_backend.is_flash: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: attn_output, attn_weights = self._kernels_flash_attn( query_states, key_states, val_states, sequence_ids=sequence_ids, is_cache_prefilled=is_cache_prefilled, ) else: raise ValueError( "E1 global attention does not support a kernels Flash backend; " "use eager, sdpa, or flex_attention." ) elif effective_backend == AttentionBackend.FLEX: attn_output, attn_weights = self._flex_attn( query_states, key_states, val_states, sequence_ids=sequence_ids, attention_args=attention_args, effective_layer_type=effective_layer_type, is_cache_prefilled=is_cache_prefilled, ) elif effective_backend == AttentionBackend.SDPA: attn_output, attn_weights = self._sdpa_attn( query_states, key_states, val_states, sequence_ids=sequence_ids, attention_args=attention_args, effective_layer_type=effective_layer_type, is_cache_prefilled=is_cache_prefilled, ) else: raise AssertionError(f"Unsupported resolved backend: {effective_backend}") s_max_key_states = key_states if ( is_cache_prefilled and effective_layer_type == AttentionLayerType.WITHIN_SEQ and query_states.shape[1] < key_states.shape[1] ): s_max_key_states = key_states[:, -query_states.shape[1] :] s_max = self._compute_s_max(query_states, s_max_key_states) if output_s_max else None return attn_output, attn_weights, s_max @torch.no_grad() def _compute_s_max( self, query_states: torch.Tensor, # Q has shape (b, l, h, d). key_states: torch.Tensor, # K has shape (b, l, h_kv, d). ) -> list[torch.Tensor]: query_heads = query_states.transpose(1, 2).contiguous() # (b, h, l, d_h) key_heads = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) key_heads = repeat_kv(key_heads, self.num_key_value_groups) scale = 1.0 / (self.head_dim**0.5) q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l) k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l) s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max( dim=0 ).values * scale return [s_max_bound[h] for h in range(self.num_heads)] def _kernels_flash_attn( self, query_states: torch.Tensor, key_states: torch.Tensor, val_states: torch.Tensor, sequence_ids: torch.Tensor, is_cache_prefilled: bool = False, ) -> tuple[torch.Tensor, None]: bsz, q_len = query_states.shape[0], query_states.shape[1] _, kv_len = key_states.shape[0], key_states.shape[1] if self.layer_type == AttentionLayerType.GLOBAL: q_sequence_ids = sequence_ids if q_len < kv_len: first_token_id = sequence_ids[:, 0].unsqueeze(1) k_sequence_ids = torch.cat( [first_token_id.expand(bsz, kv_len - q_len), sequence_ids], dim=-1 ) else: k_sequence_ids = sequence_ids else: if q_len < kv_len: key_states = key_states[:, -q_len:] val_states = val_states[:, -q_len:] q_sequence_ids = k_sequence_ids = sequence_ids attn_output = kernels_flash_attention_func( query_states, key_states, val_states, q_sequence_ids=q_sequence_ids, k_sequence_ids=k_sequence_ids, causal=False, implementation=self.attn_backend.value, ) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() return attn_output, None def _flex_attn( self, query_states: torch.Tensor, key_states: torch.Tensor, val_states: torch.Tensor, sequence_ids: torch.Tensor, attention_args: AttentionArgs | None = None, effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, is_cache_prefilled: bool = False, ) -> tuple[torch.Tensor, None]: bsz, q_len = query_states.shape[0], query_states.shape[1] kv_len = key_states.shape[1] if is_cache_prefilled and q_len < kv_len: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: key_states = key_states[:, -q_len:] val_states = val_states[:, -q_len:] block_mask = create_within_seq_block_mask(sequence_ids) outputs = flex_attention_func( query_states, key_states, val_states, block_mask=block_mask, mask_semantics=effective_layer_type.value, ) else: q_sequence_ids, k_sequence_ids = self._cached_global_sequence_ids( sequence_ids, kv_len, ) outputs = varlen_flex_attention_func( query_states, key_states, val_states, q_sequence_ids=q_sequence_ids, k_sequence_ids=k_sequence_ids, ) outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous() return outputs, None if effective_layer_type == AttentionLayerType.WITHIN_SEQ: block_mask = ( attention_args["within_seq_block_mask"] if attention_args is not None else None ) else: block_mask = ( attention_args["block_causal_block_mask"] if attention_args is not None else None ) outputs = flex_attention_func( query_states, key_states, val_states, block_mask=block_mask, mask_semantics=effective_layer_type.value, ) outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous() return outputs, None @staticmethod def _cached_global_sequence_ids( query_sequence_ids: torch.Tensor, kv_len: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Assign cached context to the incoming query sequence. E1 retrieval cache hits contain one incoming sequence. The pinned implementation relabels the cached prefix with that sequence ID so its valid query tokens attend the complete cached context, while padding is excluded by the equality mask or packed Flex path. """ q_len = query_sequence_ids.shape[1] cached_len = kv_len - q_len if cached_len < 0: raise ValueError(f"E1 cached KV length {kv_len} is shorter than query length {q_len}.") first_sequence_id = query_sequence_ids[:, :1] if bool(first_sequence_id.eq(-1).any()): raise ValueError("E1 cached queries must start with a non-padding sequence token.") cached_sequence_ids = first_sequence_id.expand(-1, cached_len) key_sequence_ids = torch.cat((cached_sequence_ids, query_sequence_ids), dim=-1) return query_sequence_ids, key_sequence_ids def _cached_attention_mask_4d( self, sequence_ids: torch.Tensor, kv_len: int, effective_layer_type: AttentionLayerType, ) -> torch.Tensor: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: return build_within_seq_mask_4d(sequence_ids) query_sequence_ids, key_sequence_ids = self._cached_global_sequence_ids( sequence_ids, kv_len, ) query_valid = query_sequence_ids.ne(-1) key_valid = key_sequence_ids.ne(-1) same_sequence = query_sequence_ids.unsqueeze(-1).eq(key_sequence_ids.unsqueeze(-2)) return (same_sequence & query_valid.unsqueeze(-1) & key_valid.unsqueeze(-2)).unsqueeze(1) def _sdpa_attn( self, query_states: torch.Tensor, # Q has shape (b, l, h, d). key_states: torch.Tensor, # K has shape (b, l, h_kv, d). val_states: torch.Tensor, # V has shape (b, l, h_kv, d). sequence_ids: torch.Tensor, attention_args: AttentionArgs | None = None, effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, is_cache_prefilled: bool = False, ) -> tuple[torch.Tensor, None]: bsz, q_len = query_states.shape[:2] kv_len = key_states.shape[1] if is_cache_prefilled and q_len < kv_len: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: key_states = key_states[:, -q_len:] val_states = val_states[:, -q_len:] attention_mask_4d = self._cached_attention_mask_4d( sequence_ids, kv_len, effective_layer_type, ) elif attention_args is not None: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: attention_mask_4d = attention_args["within_seq_mask_4d"] else: attention_mask_4d = attention_args["block_causal_mask_4d"] else: attention_mask_4d = None query_heads = query_states.transpose(1, 2).contiguous() # (b, h, l, d_h) key_heads = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) value_heads = val_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) key_heads = repeat_kv(key_heads, self.num_key_value_groups) value_heads = repeat_kv(value_heads, self.num_key_value_groups) context_heads = F.scaled_dot_product_attention( query_heads, key_heads, value_heads, attn_mask=attention_mask_4d ) # (b, h, l, d_h) attn_output = ( context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous() ) return attn_output, None def _manual_attn( self, query_states: torch.Tensor, # Q has shape (b, l, h, d). key_states: torch.Tensor, # K has shape (b, l, h_kv, d). val_states: torch.Tensor, # V has shape (b, l, h_kv, d). sequence_ids: torch.Tensor, attention_args: AttentionArgs | None = None, effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ, output_s_max: bool = False, is_cache_prefilled: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: bsz, q_len = query_states.shape[:2] kv_len = key_states.shape[1] if is_cache_prefilled and q_len < kv_len: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: key_states = key_states[:, -q_len:] val_states = val_states[:, -q_len:] attention_mask_4d = self._cached_attention_mask_4d( sequence_ids, kv_len, effective_layer_type, ) elif attention_args is not None: if effective_layer_type == AttentionLayerType.WITHIN_SEQ: attention_mask_4d = attention_args["within_seq_mask_4d"] else: attention_mask_4d = attention_args["block_causal_mask_4d"] else: attention_mask_4d = None query_heads = query_states.transpose(1, 2).contiguous() # (b, h, l, d_h) key_heads = key_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) value_heads = val_states.transpose(1, 2).contiguous() # (b, h_kv, l, d_h) key_heads = repeat_kv(key_heads, self.num_key_value_groups) value_heads = repeat_kv(value_heads, self.num_key_value_groups) scale = 1.0 / (self.head_dim**0.5) attn_weights = ( torch.matmul(query_heads, key_heads.transpose(-2, -1)) * scale ) # (b, h, l, l) if attention_mask_4d is not None: attention_mask_4d = attention_mask_4d.to(dtype=torch.bool) attn_weights = attn_weights.masked_fill( attention_mask_4d.logical_not(), torch.finfo(attn_weights.dtype).min, ) attn_weights = F.softmax(attn_weights, dim=-1) if attention_mask_4d is not None: attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), 0.0) context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) attn_output = ( context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous() ) s_max = self._compute_s_max(query_states, key_states) if output_s_max else None return attn_output, attn_weights, s_max class MLP(nn.Module): def __init__(self, config: E1Config) -> None: super().__init__() self.ffn_dim = config.intermediate_size self.hidden_dim = config.hidden_size self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.w2(self.act_fn(self.w1(hidden_states))) class GLUMLP(nn.Module): def __init__(self, config: E1Config) -> None: super().__init__() self.ffn_dim = config.intermediate_size self.hidden_dim = config.hidden_size self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False) self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False) self.act_fn = ACT2FN[config.hidden_act] def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states) hidden_states = self.w2(hidden_states) return hidden_states class FFN(nn.Module): def __init__(self, config: E1Config) -> None: super().__init__() mlp_cls = GLUMLP if config.gated_mlp else MLP self.mlp = mlp_cls(config) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.mlp(hidden_states) @dataclass class E1ModelOutputWithPast(ModelOutput): """E1 encoder outputs. ``last_hidden_state`` is H with shape (b, l, d). Optional hidden states use the same shape per layer, while attention tensors have shape (b, h, l, l). ``past_key_values`` stores the reusable K and V tensors for cached decoding. """ last_hidden_state: torch.FloatTensor | None = None past_key_values: DynamicCache | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None s_max: tuple[list[torch.Tensor], ...] | None = None @dataclass class E1MaskedLMOutputWithPast(ModelOutput): """Masked-LM output with the standard HF fields first, then E1 diagnostics.""" loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None mlm_loss: torch.FloatTensor | None = None last_hidden_state: torch.FloatTensor | None = None past_key_values: DynamicCache | None = None s_max: tuple[list[torch.Tensor], ...] | None = None @dataclass class E1ClassificationOutputWithPast(ModelOutput): """Sequence-classifier output matching HF ``SequenceClassifierOutputWithPast``.""" loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None past_key_values: DynamicCache | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None last_hidden_state: torch.FloatTensor | None = None s_max: tuple[list[torch.Tensor], ...] | None = None @dataclass class E1TokenClassificationOutputWithPast(ModelOutput): """Token-classifier output with the standard HF fields before E1 extensions.""" loss: torch.FloatTensor | None = None logits: torch.FloatTensor | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None last_hidden_state: torch.FloatTensor | None = None past_key_values: DynamicCache | None = None s_max: tuple[list[torch.Tensor], ...] | None = None class RMSNorm(nn.Module): def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps self.hidden_size = hidden_size def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype return torch.nn.functional.rms_norm( hidden_states, (self.hidden_size,), self.weight, self.variance_epsilon ).to(input_dtype) class NormAttentionNorm(nn.Module): def __init__(self, config: E1Config, layer_idx: int) -> None: super().__init__() self.self_attn = Attention(config, layer_idx) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, within_seq_position_ids: torch.LongTensor, global_position_ids: torch.LongTensor, sequence_ids: torch.LongTensor, attention_args: AttentionArgs | None = None, past_key_value: DynamicCache | None = None, output_attentions: bool = False, output_s_max: bool = False, use_cache: bool = False, effective_backend: AttentionBackend | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None, ]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value, s_max = self.self_attn( hidden_states=hidden_states, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, attention_args=attention_args, past_key_value=past_key_value, output_attentions=output_attentions, output_s_max=output_s_max, use_cache=use_cache, effective_backend=effective_backend, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) return hidden_states, residual, self_attn_weights, present_key_value, s_max class DecoderLayer(nn.Module): def __init__(self, config: E1Config, layer_idx: int) -> None: super().__init__() self.initializer_range = config.initializer_range self.hidden_size = config.hidden_size self.norm_attn_norm = NormAttentionNorm(config, layer_idx) self.ffn = FFN(config) def forward( self, hidden_states: torch.Tensor, within_seq_position_ids: torch.LongTensor, global_position_ids: torch.LongTensor, sequence_ids: torch.LongTensor, attention_args: AttentionArgs | None = None, past_key_value: DynamicCache | None = None, output_attentions: bool = False, output_s_max: bool = False, use_cache: bool = False, effective_backend: AttentionBackend | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]: hidden_states, residual, self_attn_weights, present_key_value, s_max = self.norm_attn_norm( hidden_states=hidden_states, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, attention_args=attention_args, past_key_value=past_key_value, output_attentions=output_attentions, output_s_max=output_s_max, use_cache=use_cache, effective_backend=effective_backend, ) # Fully Connected hidden_states = self.ffn(hidden_states) hidden_states = residual + hidden_states return hidden_states, self_attn_weights, present_key_value, s_max class E1PreTrainedModel(FastPLMsAttentionMixin, PreTrainedModel): config_class = E1Config embedding_unsupported_pooling = ("cls", "parti") config: E1Config base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules: ClassVar[list[str]] = ["DecoderLayer"] _transformer_layer_cls: ClassVar[list[type[nn.Module]]] = [DecoderLayer] _skip_keys_device_placement = "past_key_values" all_tied_weights_keys: ClassVar[dict[str, str]] = {} _supports_flash_attn_2 = False _supports_flash_attn_3 = False _fastplms_attention_implementations = ("sdpa", "flex_attention") _is_internal_encoder = False def __init__(self, config: E1Config, *args: Any, **kwargs: Any) -> None: super().__init__(config, *args, **kwargs) # The E1 agreement requires this exact attribution when an E1 model is # launched. Internal encoder construction is excluded so each public # model launch displays the attribution exactly once. if not self._is_internal_encoder: print("Profluent-E1", file=sys.stderr, flush=True) @classmethod def from_pretrained( # type: ignore[override] cls, pretrained_model_name_or_path: str | os.PathLike, *model_args: Any, **kwargs: Any, ) -> E1PreTrainedModel: tokenizer_token = None if "token" in kwargs: tokenizer_token = kwargs["token"] elif "use_auth_token" in kwargs: tokenizer_token = kwargs["use_auth_token"] load_context_token = _TOKENIZER_LOAD_CONTEXT.set( { "tokenizer_source": pretrained_model_name_or_path, "local_files_only": bool(kwargs.get("local_files_only", False)), "cache_dir": kwargs.get("cache_dir"), "revision": kwargs.get("revision"), "token": tokenizer_token, } ) try: return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) finally: _TOKENIZER_LOAD_CONTEXT.reset(load_context_token) @staticmethod def _tokenizer_kwargs_from_config(config: E1Config) -> dict[str, Any]: load_context = _TOKENIZER_LOAD_CONTEXT.get() resolved_revision = getattr(config, "_commit_hash", None) if not isinstance(resolved_revision, str) or not resolved_revision.strip(): resolved_revision = None if load_context is not None: tokenizer_kwargs = dict(load_context) if resolved_revision is not None: tokenizer_kwargs["revision"] = resolved_revision return tokenizer_kwargs tokenizer_source = None if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0: tokenizer_source = config._name_or_path return { "tokenizer_source": tokenizer_source, "local_files_only": False, "cache_dir": None, "revision": resolved_revision, "token": None, } @property def prep_tokens(self) -> E1BatchPreparer: """Create E1's raw-sequence preparer only when a sequence API uses it.""" preparer = self.__dict__.get("_fastplms_prep_tokens") if preparer is not None: return preparer encoder = self._modules.get("model") if encoder is not None and encoder is not self: return encoder.prep_tokens tokenizer_kwargs = self.__dict__.get("_fastplms_tokenizer_kwargs") if tokenizer_kwargs is None: raise RuntimeError("E1 tokenizer settings were not initialized.") preparer = E1BatchPreparer( data_prep_config=DataPrepConfig( max_num_sequences=self.config.max_num_sequences, max_num_positions_within_seq=self.config.max_num_positions_within_seq, ), **tokenizer_kwargs, ) self.__dict__["_fastplms_prep_tokens"] = preparer return preparer @prep_tokens.setter def prep_tokens(self, value: E1BatchPreparer | None) -> None: self.__dict__["_fastplms_prep_tokens"] = value def _init_weights(self, module: nn.Module) -> None: if isinstance(module, RMSNorm): nn.init.ones_(module.weight) return if not isinstance(module, (nn.Linear, nn.Embedding)): return nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) if isinstance(module, nn.Linear) and module.bias is not None: nn.init.zeros_(module.bias) if isinstance(module, nn.Embedding) and module.padding_idx is not None: with torch.no_grad(): module.weight[module.padding_idx].zero_() def _backward_compatibility_gradient_checkpointing(self) -> None: if self.supports_gradient_checkpointing and getattr( self.config, "gradient_checkpointing", False ): self.gradient_checkpointing_enable(dict(use_reentrant=False)) def post_init(self) -> None: super().post_init() @property def _device(self) -> torch.device: return next(self.parameters()).device @property def attn_backend(self) -> str: return self.config.attn_backend @attn_backend.setter def attn_backend(self, backend: str) -> None: if backend not in self._fastplms_attention_implementations: raise ValueError( f"E1 does not support {backend!r}; expected one of " f"{self._fastplms_attention_implementations}." ) self.config.attn_backend = backend resolved = resolve_attention_backend(backend) for module in self.modules(): if isinstance(module, FAST_E1_ENCODER): module._attn_backend = resolved elif isinstance(module, Attention): module.attn_backend = resolved class FAST_E1_ENCODER(E1PreTrainedModel, EmbeddingMixin): config: E1Config config_class = E1Config _is_internal_encoder = True def __init__(self, config: E1Config, **kwargs) -> None: E1PreTrainedModel.__init__(self, config, **kwargs) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.embed_seq_id = nn.Embedding(config.max_num_sequences, config.hidden_size) self.layers = nn.ModuleList( [DecoderLayer(config, i) for i in range(config.num_hidden_layers)] ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.gradient_checkpointing = config.gradient_checkpointing self.__dict__["_fastplms_tokenizer_kwargs"] = ( E1PreTrainedModel._tokenizer_kwargs_from_config(config) ) self.__dict__["_fastplms_prep_tokens"] = None self._attn_backend = resolve_attention_backend(config.attn_backend) self.post_init() def get_input_embeddings(self) -> nn.Embedding: return self.embed_tokens def set_input_embeddings(self, value: nn.Embedding) -> None: self.embed_tokens = value def _embed( self, sequences: list[str], return_attention_mask: bool = False, hidden_state_index: int = -1, store_all_hidden_states: bool = False, **kwargs, ) -> torch.Tensor: batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device) # The native preparer also returns training labels plus retrieval # descriptors. The encoder accepts only its aligned model inputs. encoder_batch: dict[str, torch.Tensor] = {} for name in ( "input_ids", "within_seq_position_ids", "global_position_ids", "sequence_ids", ): value = batch[name] if not isinstance(value, torch.Tensor): raise TypeError(f"Prepared E1 field {name!r} must be a tensor.") encoder_batch[name] = value output_hidden_states = store_all_hidden_states or hidden_state_index != -1 output = self.forward( **encoder_batch, output_hidden_states=output_hidden_states, output_attentions=False, return_dict=True, ) embeddings = select_hidden_state_embeddings( output.last_hidden_state, output.hidden_states, hidden_state_index=hidden_state_index, store_all_hidden_states=store_all_hidden_states, ) if return_attention_mask: attention_mask = (encoder_batch["sequence_ids"] != -1).long() return embeddings, attention_mask else: return embeddings def _prepare_hidden_states( self, input_ids: torch.LongTensor | None, inputs_embeds: torch.FloatTensor | None, within_seq_position_ids: torch.LongTensor | None, global_position_ids: torch.LongTensor | None, sequence_ids: torch.LongTensor | None, ) -> tuple[torch.Tensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]: if (input_ids is None) == (inputs_embeds is None): message = ( "Must specify either input_ids or inputs_embeds" if input_ids is None else "Cannot specify both input_ids and inputs_embeds" ) raise ValueError(message) source = input_ids if input_ids is not None else inputs_embeds if source is None: raise RuntimeError("E1 input validation did not resolve an input tensor.") expected_rank = 2 if input_ids is not None else 3 if source.ndim != expected_rank: source_name = "input_ids" if input_ids is not None else "inputs_embeds" raise ValueError( f"{source_name} must have rank {expected_rank}; got shape {tuple(source.shape)}." ) batch_size, sequence_length = source.shape[:2] if sequence_length == 0: raise ValueError("E1 inputs must contain at least one token.") if inputs_embeds is not None and inputs_embeds.shape[-1] != self.config.hidden_size: raise ValueError( "inputs_embeds hidden dimension must match config.hidden_size; " f"got {inputs_embeds.shape[-1]} and {self.config.hidden_size}." ) if inputs_embeds is not None: default_positions = torch.arange(sequence_length, device=source.device).expand( batch_size, -1, ) if within_seq_position_ids is None: within_seq_position_ids = default_positions if global_position_ids is None: global_position_ids = default_positions if sequence_ids is None: sequence_ids = torch.zeros_like(default_positions) if within_seq_position_ids is None or global_position_ids is None or sequence_ids is None: raise ValueError("Position and sequence IDs are required when input_ids are provided.") expected_shape = (batch_size, sequence_length) aligned_inputs = { "within_seq_position_ids": within_seq_position_ids, "global_position_ids": global_position_ids, "sequence_ids": sequence_ids, } for name, value in aligned_inputs.items(): if tuple(value.shape) != expected_shape: raise ValueError( f"{name} must have shape {expected_shape}; got {tuple(value.shape)}." ) within_positions = within_seq_position_ids.long() global_positions = global_position_ids.long() sequence_numbers = sequence_ids.long() lowest_position, highest_position = torch.aminmax(within_positions) if ( lowest_position.item() < -1 or highest_position.item() >= self.config.max_num_positions_within_seq ): raise ValueError( "Position ids must be in the range " f"[-1, {self.config.max_num_positions_within_seq}); got max " f"{highest_position.item()} and min {lowest_position.item()}" ) lowest_global, highest_global = torch.aminmax(global_positions) if ( lowest_global.item() < -1 or highest_global.item() >= self.config.max_num_positions_global ): raise ValueError( "Global position ids must be in the range " f"[-1, {self.config.max_num_positions_global}); got max " f"{highest_global.item()} and min {lowest_global.item()}" ) lowest_sequence, highest_sequence = torch.aminmax(sequence_numbers) if lowest_sequence.item() < -1 or highest_sequence.item() >= self.config.max_num_sequences: raise ValueError( "Sequence ids must be in the range " f"[-1, {self.config.max_num_sequences}); got max " f"{highest_sequence.item()} and min {lowest_sequence.item()}" ) if inputs_embeds is None: if input_ids is None: raise RuntimeError("E1 input validation lost the token ID tensor.") token_embeddings = self.embed_tokens(input_ids) inputs_embeds = token_embeddings + self.embed_seq_id(sequence_numbers.clamp_min(0)) layer_dtype = self.layers[0].norm_attn_norm.self_attn.q_proj.weight.dtype target_dtype = ( torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else layer_dtype ) return ( inputs_embeds.to(target_dtype), within_positions, global_positions, sequence_numbers, ) def _resolve_forward_cache( self, past_key_values: DynamicCache | None, use_cache: bool, ) -> tuple[DynamicCache | None, bool]: checkpointing = self.gradient_checkpointing and self.training and torch.is_grad_enabled() if checkpointing and use_cache: _get_logger().warning_once( "`use_cache=True` is incompatible with gradient checkpointing; " "setting `use_cache=False`." ) use_cache = False if not use_cache: return None, False return past_key_values if past_key_values is not None else DynamicCache(), True def _build_forward_attention_args( self, sequence_ids: torch.LongTensor, past_key_values: DynamicCache | None, effective_backend: AttentionBackend, ) -> AttentionArgs | None: if past_key_values is not None and past_key_values.get_seq_length() != 0: return None use_flex = effective_backend == AttentionBackend.FLEX use_dense_mask = effective_backend in { AttentionBackend.EAGER, AttentionBackend.SDPA, } return AttentionArgs( block_causal_block_mask=( create_block_causal_mask_optimized(sequence_ids) if use_flex and self.config.global_attention_every_n_layers > 0 else None ), within_seq_block_mask=( create_within_seq_block_mask(sequence_ids) if use_flex else None ), within_seq_mask_4d=(build_within_seq_mask_4d(sequence_ids) if use_dense_mask else None), block_causal_mask_4d=( build_block_causal_mask_4d(sequence_ids) if use_dense_mask else None ), ) def _run_decoder_layers( self, hidden_states: torch.Tensor, within_seq_position_ids: torch.LongTensor, global_position_ids: torch.LongTensor, sequence_ids: torch.LongTensor, attention_args: AttentionArgs | None, past_key_values: DynamicCache | None, use_cache: bool, output_attentions: bool, output_hidden_states: bool, output_s_max: bool, effective_backend: AttentionBackend, ) -> E1ModelOutputWithPast: hidden_history: list[torch.Tensor] | None = [] if output_hidden_states else None attention_history: list[torch.Tensor] | None = [] if output_attentions else None s_max_history: list[list[torch.Tensor]] | None = [] if output_s_max else None next_cache: DynamicCache | None = None for layer in self.layers: if hidden_history is not None: hidden_history.append(hidden_states) if self.gradient_checkpointing and self.training and torch.is_grad_enabled(): layer_output = self._gradient_checkpointing_func( layer.__call__, hidden_states, within_seq_position_ids, global_position_ids, sequence_ids, attention_args, past_key_values, output_attentions, output_s_max, use_cache, effective_backend, ) else: layer_output = layer( hidden_states, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, attention_args=attention_args, past_key_value=past_key_values, output_attentions=output_attentions, output_s_max=output_s_max, use_cache=use_cache, effective_backend=effective_backend, ) hidden_states, attention, layer_cache, s_max = layer_output if use_cache: past_key_values = layer_cache next_cache = layer_cache if attention_history is not None: if attention is None: raise RuntimeError( "An E1 layer did not return attention tensors when requested." ) attention_history.append(attention) if s_max_history is not None: if s_max is None: raise RuntimeError( "An E1 layer did not return s_max diagnostics when requested." ) s_max_history.append(s_max) hidden_states = self.norm(hidden_states) if hidden_history is not None: hidden_history.append(hidden_states) return E1ModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=tuple(hidden_history) if hidden_history is not None else None, attentions=tuple(attention_history) if attention_history is not None else None, s_max=tuple(s_max_history) if s_max_history is not None else None, ) def forward( self, input_ids: torch.LongTensor | None = None, within_seq_position_ids: torch.LongTensor | None = None, global_position_ids: torch.LongTensor | None = None, sequence_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, past_key_values: DynamicCache | None = None, use_cache: bool | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, output_s_max: bool = False, return_dict: bool | None = None, ) -> E1ModelOutputWithPast | tuple[Any, ...]: """Transform token or soft embeddings H with shape (b, l, d).""" use_cache = ( use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) ) output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict hidden_states, within_positions, global_positions, sequence_numbers = ( self._prepare_hidden_states( input_ids, inputs_embeds, within_seq_position_ids, global_position_ids, sequence_ids, ) ) cache, use_cache = self._resolve_forward_cache(past_key_values, use_cache) effective_backend = resolve_attention_backend_for_call( self._attn_backend, output_attentions=bool(output_attentions), ) attention_args = self._build_forward_attention_args( sequence_numbers, cache, effective_backend, ) result = self._run_decoder_layers( hidden_states, within_positions, global_positions, sequence_numbers, attention_args, cache, use_cache, output_attentions, output_hidden_states, output_s_max, effective_backend, ) if not return_dict: return result.to_tuple() return result class E1Model(E1PreTrainedModel, EmbeddingMixin): config: E1Config config_class = E1Config def __init__(self, config: E1Config, **kwargs) -> None: E1PreTrainedModel.__init__(self, config, **kwargs) self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) self.post_init() def get_input_embeddings(self) -> nn.Embedding: return self.model.get_input_embeddings() def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.set_input_embeddings(value) def _embed( self, sequences: list[str], return_attention_mask: bool = False, **kwargs ) -> torch.Tensor: return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) def forward( self, input_ids: torch.LongTensor | None = None, within_seq_position_ids: torch.LongTensor | None = None, global_position_ids: torch.LongTensor | None = None, sequence_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, past_key_values: DynamicCache | None = None, use_cache: bool | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, output_s_max: bool = False, return_dict: bool | None = None, ) -> E1ModelOutputWithPast | tuple[Any, ...]: return self.model( input_ids=input_ids, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_s_max=output_s_max, return_dict=return_dict, ) class E1ForMaskedLM(FastPLMTestTimeTrainingMixin, E1PreTrainedModel, EmbeddingMixin): config: E1Config config_class = E1Config def __init__(self, config: E1Config, **kwargs) -> None: E1PreTrainedModel.__init__(self, config, **kwargs) self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) self.vocab_size = config.vocab_size self.mlm_head = torch.nn.Sequential( nn.Linear(config.hidden_size, config.hidden_size, bias=True), nn.GELU(), nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps), nn.Linear(config.hidden_size, config.vocab_size, bias=True), ) self.gradient_checkpointing = config.gradient_checkpointing self.post_init() self.init_ttt({"lora_target_replace_module": "Attention"}) @property def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: return self.model.device_mesh def get_input_embeddings(self) -> nn.Embedding: return self.model.get_input_embeddings() def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.set_input_embeddings(value) def _embed( self, sequences: list[str], return_attention_mask: bool = False, **kwargs ) -> torch.Tensor: return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) def get_output_embeddings(self) -> nn.Linear: return self.mlm_head[-1] def set_output_embeddings(self, value: nn.Linear) -> None: self.mlm_head[-1] = value def _ttt_get_trainable_modules(self) -> list[nn.Module]: return [self.model] def _ttt_tokenize( self, seq: str | list[str] | None = None, input_ids: torch.Tensor | None = None, **kwargs, ) -> dict[str, torch.Tensor]: if input_ids is not None: return { "input_ids": input_ids, "within_seq_position_ids": kwargs["within_seq_position_ids"], "global_position_ids": kwargs["global_position_ids"], "sequence_ids": kwargs["sequence_ids"], } if seq is None: raise ValueError("Pass either seq or E1 token tensors for TTT.") sequences = [seq] if isinstance(seq, str) else seq batch = self.prep_tokens.get_batch_kwargs(sequences, device=torch.device("cpu")) return { "input_ids": batch["input_ids"], "within_seq_position_ids": batch["within_seq_position_ids"], "global_position_ids": batch["global_position_ids"], "sequence_ids": batch["sequence_ids"], } def _ttt_mask_token(self) -> int: return int(self.prep_tokens.mask_token_id) def _ttt_padding_token(self) -> int: return int(self.prep_tokens.pad_token_id) def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: amino_acids = "ACDEFGHIKLMNPQRSTVWY" ids = [self.prep_tokens.vocab[aa] for aa in amino_acids] return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype) def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: return ~self.prep_tokens.get_boundary_token_mask(input_ids) def _ttt_predict_logits( self, batch: torch.Tensor | dict[str, torch.Tensor], **kwargs, ) -> torch.Tensor: del kwargs if not isinstance(batch, dict): raise TypeError("E1 TTT expects a tensor dictionary.") output = self( input_ids=batch["input_ids"], within_seq_position_ids=batch["within_seq_position_ids"], global_position_ids=batch["global_position_ids"], sequence_ids=batch["sequence_ids"], return_dict=True, ) return output.logits def search_homologues( self, sequence: str, output_dir: str, provider: str = "colabfold", target_db: str | None = None, seq_id: str | None = None, **kwargs, ) -> str: searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs) return searcher.search(sequence=sequence, output_dir=output_dir, seq_id=seq_id) def batch_search_homologues( self, sequences: list[str], output_dir: str, provider: str = "colabfold", target_db: str | None = None, seq_ids: list[str] | None = None, continue_on_error: bool = True, **kwargs, ) -> dict[str, str]: searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs) return searcher.batch_search( sequences=sequences, output_dir=output_dir, seq_ids=seq_ids, continue_on_error=continue_on_error, ) def sample_msa_contexts( self, a3m_path: str, seed: int = 42, max_context_tokens: list[int] | None = None, similarity_thresholds: list[float] | None = None, min_query_similarity: float = 0.3, context_cache_dir: str | None = None, ) -> dict[str, str]: context_specs = build_context_specifications( max_context_tokens=max_context_tokens, similarity_thresholds=similarity_thresholds, min_query_similarity=min_query_similarity, ) cache = None if context_cache_dir is not None: key = repr((max_context_tokens, similarity_thresholds, min_query_similarity)) specs_hash = hashlib.md5(key.encode()).hexdigest()[:8] cache = ContextCache(context_cache_dir, specs_hash, seed) cached = cache.load(a3m_path) if cached is not None: return cached contexts = sample_contexts_for_msa(a3m_path, context_specs, seed=seed) if cache is not None: cache.store(a3m_path, contexts) return contexts @torch.inference_mode() def score_ppll( self, sequences: list[str], a3m_path: str, ensemble: bool = True, seed: int = 42, max_context_tokens: list[int] | None = None, similarity_thresholds: list[float] | None = None, min_query_similarity: float = 0.3, max_batch_tokens: int = 131072, cache_size: int = 1, context_cache_dir: str | None = None, progress: bool = True, ) -> list[float] | list[list[float]]: """Score sequences with FastPLMs PPLL reduction over sampled E1 MSA contexts. This intentionally differs from Profluent's official E1Scorer, which scores mutants against a parent sequence with wildtype or masked marginal log-prob deltas. Here each sequence is scored by mean correct-token probability and optionally averaged across sampled contexts. """ contexts = self.sample_msa_contexts( a3m_path=a3m_path, seed=seed, max_context_tokens=max_context_tokens, similarity_thresholds=similarity_thresholds, min_query_similarity=min_query_similarity, context_cache_dir=context_cache_dir, ) if not contexts: raise ValueError("At least one sampled MSA context is required for PPLL scoring.") predictor = _E1ContextPredictor( model=self, data_prep_config=DataPrepConfig(remove_X_tokens=True), max_batch_tokens=max_batch_tokens, fields_to_save=["logits"], save_masked_positions_only=False, keep_predictions_in_gpu=False, use_cache=True, cache_size=cache_size, progress=progress, ) vocab = predictor.batch_preparer.vocab seq_token_ids = [ torch.tensor([vocab[aa] for aa in seq if aa != "X"], device=self.device) for seq in sequences ] context_ids = list(contexts.keys()) all_scores = torch.zeros(len(sequences), len(context_ids), device=self.device) iterator = tqdm(context_ids, desc="Scoring with contexts", disable=not progress) for ctx_idx, ctx_id in enumerate(iterator): predictions = list( predictor.predict( sequences=sequences, sequence_ids=list(range(len(sequences))), context_seqs={ctx_id: contexts[ctx_id]}, ) ) for prediction in predictions: seq_idx = prediction["id"] if not isinstance(seq_idx, int): raise TypeError("Expected integer sequence ids for score aggregation.") all_scores[seq_idx, ctx_idx] = compute_ppll( prediction["logits"], seq_token_ids[seq_idx] ) if predictor.kv_cache is not None: predictor.kv_cache.reset() if ensemble: return all_scores.mean(dim=1).tolist() return all_scores.tolist() @torch.inference_mode() def embed_with_msa( self, sequences: list[str], a3m_path: str | None = None, context: str | None = None, pooling_types: list[str] | None = None, pooling: str = "mean", matrix_embed: bool = False, seed: int = 42, max_batch_tokens: int = 131072, embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS, embed_similarity: float = DEFAULT_EMBED_SIMILARITY, min_query_similarity: float = 0.3, progress: bool = True, ) -> torch.Tensor | list[torch.Tensor]: if a3m_path is not None and context is None: spec = ContextSpecification( max_num_samples=511, max_token_length=embed_max_tokens, max_query_similarity=embed_similarity, min_query_similarity=min_query_similarity, ) contexts, _ = sample_multiple_contexts( msa_path=a3m_path, context_specifications=[spec], seed=seed, ) context = contexts[0] if contexts else None hidden_list = _forward_for_embedding( model=self, sequences=sequences, context=context, max_batch_tokens=max_batch_tokens, progress=progress, ) if matrix_embed: return hidden_list if pooling_types is not None: return _pool_hidden_states(hidden_list, pooling_types, self.device) if pooling not in ("mean", "cls"): raise ValueError("pooling must be 'mean' or 'cls' when pooling_types is not provided") embeddings = [ hidden.mean(dim=0) if pooling == "mean" else hidden[0] for hidden in hidden_list ] return torch.stack(embeddings) @torch.inference_mode() def embed_dataset_with_msa( self, sequences: list[str], msa_lookup: dict[str, str] | None = None, msa_dir: str | None = None, msa_hf_path: str | None = None, batch_size: int = 2, max_len: int = 2048, pooling_types: list[str] | None = None, pooling: str = "mean", matrix_embed: bool = False, embed_dtype: torch.dtype = torch.bfloat16, embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS, embed_similarity: float = DEFAULT_EMBED_SIMILARITY, min_query_similarity: float = 0.3, seed: int = 42, progress: bool = True, max_batch_tokens: int = 131072, batch_window_size: int | None = None, max_tokens_per_batch: int | None = None, output: str | os.PathLike[str] | None = None, format: str = "safetensors", resume: bool = True, shard_size: int = 2 * 1024**3, model_state_fingerprint: str | None = None, ) -> EmbeddingResult: """Embed an ordered sequence dataset with optional sampled MSA context. Unlike the legacy dictionary return, the result preserves duplicate sequences and input order. ``output`` uses the same transactional, resumable SQLite or safetensors persistence as :meth:`embed_dataset`. ``max_len`` counts biological residues. """ if not sequences: raise ValueError("sequences must contain at least one protein sequence.") if any(not isinstance(sequence, str) or not sequence for sequence in sequences): raise ValueError("sequences must contain non-empty strings.") if max_len <= 0: raise ValueError("max_len must be positive.") if msa_lookup is None: if msa_dir is not None: msa_lookup = load_msa_dir(msa_dir) elif msa_hf_path is not None: msa_lookup = load_msa_from_hf(msa_hf_path) else: msa_lookup = {} truncated_sequences = [sequence[:max_len] for sequence in sequences] unique_seqs = sorted(set(truncated_sequences), key=lambda value: (-len(value), value)) context_map: dict[str, str | None] = {} spec = ContextSpecification( max_num_samples=511, max_token_length=embed_max_tokens, max_query_similarity=embed_similarity, min_query_similarity=min_query_similarity, ) for seq in unique_seqs: a3m_path = get_msa_for_sequence(seq, msa_lookup) if a3m_path is None: context_map[seq] = None continue contexts, _ = sample_multiple_contexts( msa_path=a3m_path, context_specifications=[spec], seed=seed, ) context_map[seq] = contexts[0] if contexts else None context_digest = hashlib.sha256() for sequence in unique_seqs: for value in (sequence, context_map[sequence] or ""): encoded = value.encode("utf-8") context_digest.update(len(encoded).to_bytes(8, "big")) context_digest.update(encoded) def embed_msa_batch(batch_sequences: list[str]) -> EmbeddingBatch: grouped_positions: dict[str | None, list[int]] = defaultdict(list) for position, sequence in enumerate(batch_sequences): grouped_positions[context_map[sequence]].append(position) hidden_by_position: list[torch.Tensor | None] = [None] * len(batch_sequences) for context, positions in grouped_positions.items(): context_sequences = [batch_sequences[position] for position in positions] hidden_states = _forward_for_embedding( model=self, sequences=context_sequences, context=context, max_batch_tokens=max_batch_tokens, progress=progress, ) for position, hidden in zip(positions, hidden_states, strict=True): hidden_by_position[position] = hidden resolved = [hidden for hidden in hidden_by_position if hidden is not None] if len(resolved) != len(batch_sequences): raise RuntimeError("E1 MSA embedding did not return every requested sequence.") max_residues = max(hidden.shape[0] for hidden in resolved) hidden_size = resolved[0].shape[-1] X = resolved[0].new_zeros((len(resolved), max_residues, hidden_size)) residue_mask = torch.zeros( (len(resolved), max_residues), dtype=torch.bool, device=X.device, ) for position, hidden in enumerate(resolved): residue_count = hidden.shape[0] X[position, :residue_count] = hidden residue_mask[position, :residue_count] = True return EmbeddingBatch(X=X, residue_mask=residue_mask) resolved_pooling: str | list[str] | None = ( None if matrix_embed else pooling_types if pooling_types is not None else pooling ) adapter_identity = { "kind": "e1-msa-v1", "sampling_source_revision": E1_MSA_SAMPLING_SOURCE_REVISION, "context_sha256": context_digest.hexdigest(), "context_count": sum(context is not None for context in context_map.values()), "seed": seed, "embed_max_tokens": embed_max_tokens, "embed_similarity": embed_similarity, "min_query_similarity": min_query_similarity, "max_batch_tokens": max_batch_tokens, } return embed_dataset( self, [(str(position), sequence) for position, sequence in enumerate(sequences)], batch_size=batch_size, pooling=resolved_pooling, full_embeddings=matrix_embed, output=output, format=format, resume=resume, max_length=max_len, truncate=True, dtype=embed_dtype, shard_size=shard_size, model_state_fingerprint=model_state_fingerprint, batch_window_size=batch_window_size, max_tokens_per_batch=max_tokens_per_batch, _embedding_batch_fn=embed_msa_batch, _embedding_batch_identity=adapter_identity, _allowed_unsupported_pooling=("cls",), ) def forward( self, input_ids: torch.LongTensor | None = None, within_seq_position_ids: torch.LongTensor | None = None, global_position_ids: torch.LongTensor | None = None, sequence_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, past_key_values: DynamicCache | None = None, use_cache: bool | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, output_s_max: bool = False, return_dict: bool | None = None, ) -> E1MaskedLMOutputWithPast | tuple[Any, ...]: """Return hidden states and masked-token logits for E1 inputs. Token, position, sequence, and label tensors have shape (b, l). Callers may instead provide precomputed H with shape (b, l, d). """ use_cache = ( use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) ) output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs: E1ModelOutputWithPast = self.model( input_ids=input_ids, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_s_max=output_s_max, return_dict=True, ) last_hidden_state = outputs.last_hidden_state loss = None mlm_logits = self.mlm_head(last_hidden_state).float() mlm_loss = None if labels is not None: mlm_logits_flat = mlm_logits.contiguous().view(-1, self.config.vocab_size) mlm_labels_flat = labels.to(mlm_logits_flat.device).contiguous().view(-1) mlm_loss = F.cross_entropy( mlm_logits_flat, mlm_labels_flat, ignore_index=-100, reduction="none", ) mask = mlm_labels_flat.ne(-100) & mlm_labels_flat.ne(self.model.padding_idx) n_mlm = mask.sum().clamp_min(1) mlm_loss = (mlm_loss * mask.to(mlm_loss)).sum() / n_mlm loss = 0.0 loss += mlm_loss result = E1MaskedLMOutputWithPast( loss=loss, logits=mlm_logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, mlm_loss=mlm_loss, last_hidden_state=last_hidden_state, past_key_values=outputs.past_key_values, s_max=outputs.s_max, ) if not return_dict: return result.to_tuple() return result class E1ForSequenceClassification(E1PreTrainedModel, EmbeddingMixin): config: E1Config config_class = E1Config def __init__(self, config: E1Config, **kwargs) -> None: pooling_types = kwargs.pop("pooling_types", None) if pooling_types is None: pooling_types = ["mean", "var"] elif not isinstance(pooling_types, list): raise TypeError("pooling_types must be a non-empty list of pooling names") elif not pooling_types or not all(isinstance(name, str) for name in pooling_types): raise ValueError("pooling_types must be a non-empty list of pooling names") E1PreTrainedModel.__init__(self, config, **kwargs) self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) self.vocab_size = config.vocab_size self.num_labels = config.num_labels self.pooler = Pooler(pooling_types) self.classifier = nn.Sequential( nn.Linear(config.hidden_size * len(pooling_types), config.hidden_size * 4), nn.GELU(), nn.LayerNorm(config.hidden_size * 4), nn.Linear(config.hidden_size * 4, config.num_labels), ) self.mse = nn.MSELoss() self.ce = nn.CrossEntropyLoss() self.bce = nn.BCEWithLogitsLoss() self.gradient_checkpointing = config.gradient_checkpointing self.post_init() @property def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: return self.model.device_mesh def get_input_embeddings(self) -> nn.Embedding: return self.model.get_input_embeddings() def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.set_input_embeddings(value) def _embed( self, sequences: list[str], return_attention_mask: bool = False, **kwargs ) -> torch.Tensor: return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) def forward( self, input_ids: torch.LongTensor | None = None, within_seq_position_ids: torch.LongTensor | None = None, global_position_ids: torch.LongTensor | None = None, sequence_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, past_key_values: DynamicCache | None = None, use_cache: bool | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, output_s_max: bool = False, return_dict: bool | None = None, ) -> E1ClassificationOutputWithPast | tuple[Any, ...]: use_cache = ( use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) ) output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs: E1ModelOutputWithPast = self.model( input_ids=input_ids, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_s_max=output_s_max, return_dict=True, ) attention_mask = ( (sequence_ids != -1).long() if sequence_ids is not None else torch.ones( outputs.last_hidden_state.shape[:2], device=outputs.last_hidden_state.device, dtype=torch.long, ) ) x = outputs.last_hidden_state features = self.pooler(x, attention_mask) logits = self.classifier(features) loss = None if labels is not None: labels = labels.to(logits.device) if self.config.problem_type is None: if self.num_labels == 1: self.config.problem_type = "regression" elif self.num_labels > 1 and ( labels.dtype == torch.long or labels.dtype == torch.int ): self.config.problem_type = "single_label_classification" else: self.config.problem_type = "multi_label_classification" if self.config.problem_type == "regression": if self.num_labels == 1: loss = self.mse(logits.flatten(), labels.flatten()) else: loss = self.mse(logits, labels) elif self.config.problem_type == "single_label_classification": loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) elif self.config.problem_type == "multi_label_classification": loss = self.bce(logits, labels) result = E1ClassificationOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, last_hidden_state=x, s_max=outputs.s_max, ) if not return_dict: return result.to_tuple() return result class E1ForTokenClassification(E1PreTrainedModel, EmbeddingMixin): config: E1Config config_class = E1Config def __init__(self, config: E1Config, **kwargs) -> None: E1PreTrainedModel.__init__(self, config, **kwargs) self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs) self.vocab_size = config.vocab_size self.num_labels = config.num_labels self.classifier = nn.Sequential( nn.Linear(config.hidden_size, config.hidden_size * 4), nn.GELU(), nn.LayerNorm(config.hidden_size * 4), nn.Linear(config.hidden_size * 4, config.num_labels), ) self.loss_fct = nn.CrossEntropyLoss() self.gradient_checkpointing = config.gradient_checkpointing self.post_init() @property def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh: return self.model.device_mesh def get_input_embeddings(self) -> nn.Embedding: return self.model.get_input_embeddings() def set_input_embeddings(self, value: nn.Embedding) -> None: self.model.set_input_embeddings(value) def _embed( self, sequences: list[str], return_attention_mask: bool = False, **kwargs ) -> torch.Tensor: return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs) def forward( self, input_ids: torch.LongTensor | None = None, within_seq_position_ids: torch.LongTensor | None = None, global_position_ids: torch.LongTensor | None = None, sequence_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, past_key_values: DynamicCache | None = None, use_cache: bool | None = None, output_attentions: bool | None = None, output_hidden_states: bool | None = None, output_s_max: bool = False, return_dict: bool | None = None, ) -> E1TokenClassificationOutputWithPast | tuple[Any, ...]: use_cache = ( use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False)) ) output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs: E1ModelOutputWithPast = self.model( input_ids=input_ids, within_seq_position_ids=within_seq_position_ids, global_position_ids=global_position_ids, sequence_ids=sequence_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, output_s_max=output_s_max, return_dict=True, ) x = outputs.last_hidden_state logits = self.classifier(x) loss = None if labels is not None: labels = labels.to(logits.device) loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) result = E1TokenClassificationOutputWithPast( loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions, last_hidden_state=x, past_key_values=outputs.past_key_values, s_max=outputs.s_max, ) if not return_dict: return result.to_tuple() return result