import torch import numpy as np import math import logging from typing import Any, Optional, Union, Sequence import torch.nn as nn from transformers import PreTrainedModel, T5EncoderModel, T5ForConditionalGeneration, T5ForQuestionAnswering, T5ForTokenClassification, T5Model from torch import nn from transformers.models.t5.modeling_t5 import T5Attention, T5DenseActDense, T5DenseGatedActDense, T5ClassificationHead, T5LayerNorm, T5Stack, T5Block, T5LayerSelfAttention, T5LayerFF from transformers.cache_utils import DynamicCache, EncoderDecoderCache from transformers.models.t5.configuration_t5 import T5Config from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, BaseModelOutput from transformers.utils import DUMMY_INPUTS, DUMMY_MASK, is_torch_fx_proxy, is_torchdynamo_compiling from transformers.utils.deprecation import deprecate_kwarg from .common import M5Pooler from .prepare_data import get_positional_encodings_and_align logger = logging.getLogger(__name__) class M5EncoderConfig(T5Config): model_type = "m5_model" def __init__( self, d_ff= 2048, d_kv = 64, d_model = 512, num_layers = 24, num_heads = 12, pad_token_id = 2, dropout_rate = 0, feed_forward_proj = "gated-gelu", classifier_dropout=0, relative_attention_max_distance=96, relative_attention_num_buckets=32, vocab_size=1032, num_decoder_layers=0, **kwargs, ): super().__init__(d_ff=d_ff, d_kv=d_kv, d_model=d_model, num_layers=num_layers, num_heads=num_heads, pad_token_id=pad_token_id, dropout_rate=dropout_rate, feed_forward_proj=feed_forward_proj, classifier_dropout=classifier_dropout, relative_attention_max_distance=relative_attention_max_distance, relative_attention_num_buckets=relative_attention_num_buckets, vocab_size=vocab_size, num_decoder_layers=num_decoder_layers, **kwargs) class M5Encoder(PreTrainedModel): config_class = M5EncoderConfig base_model_prefix = "encoder" def __init__(self, config): super().__init__(config) self.model = M5EncoderModel(config) def get_input_embeddings(self): return self.model.shared def set_input_embeddings(self, new_embeddings): self.model.shared = new_embeddings self.model.encoder.embed_tokens = new_embeddings # keep encoder in sync def forward(self, input_ids, attention_mask=None, relative_position=None, **kwargs): return self.model(input_ids=input_ids, attention_mask=attention_mask, relative_position=relative_position) @staticmethod def get_positional_encodings_and_align( smiles: str, seed: int, token_regr: Optional[np.ndarray] = None, ) -> tuple[str, np.ndarray, Optional[np.ndarray]]: """ Convert a SMILES string into a SELFIES tokenization, compute pairwise molecular-graph distance encodings, and optionally align token-level regression labels to the new token order. Args: smiles: Input molecule as a SMILES string. Does not need to be canonical — canonicalization and optional randomization are applied internally. seed: Epoch/seed value controlling SMILES augmentation. When 0, the canonical SELFIES is used; any other value produces a reproducible randomized SELFIES variant. token_regr: Optional array for reproducibility. Array of per-atom regression labels (e.g. Löwdin charges) aligned to the original SMILES atom order. If provided, labels are re-aligned to match the SELFIES token order of the (possibly randomized) output SMILES. Shape: ``(n_atoms,)``. Returns: A tuple of: - **selfies** (``str``): SELFIES encoding of the (possibly randomized) SMILES. - **pos_encod** (``np.ndarray``): Pairwise distance matrix of shape ``(seq_len, seq_len)`` with ``dtype=np.int16``. Entries are shortest-path graph distances between atoms, capped at ``np.iinfo(np.int16).max - 1``. Special values: ``0`` for CLS-to-token, token-to-CLS, and ring/dot-separated fragment pairs; ``-1`` for intra-branch/ring structural tokens; ``np.iinfo(np.int16).max`` for padding positions. - **token_regr_selfies** (``np.ndarray`` or ``None``): Labels re-aligned to SELFIES token positions, shape ``(seq_len - 1,)``, with ``np.nan`` for non-atom tokens (branches, rings, dots). ``None`` if ``token_regr`` was not provided. """ return get_positional_encodings_and_align(smiles, token_regr, seed) @staticmethod def collate_for_dataset(batch: list[dict[str, Any]], n_global_regr: int = 0, PAD_TOKEN_ID: int = 2): """ Collate processed data for pytorch dataloaders. Each item in ``batch`` is a 3-tuple ``(token_dict, pos_encod, reg)`` where: - ``token_dict`` is a dict with keys ``"input_ids"`` (``np.ndarray``, shape ``(L,)``) and ``"attention_mask"`` (``np.ndarray``, shape ``(L,)``), as produced by a tokenizer. - ``pos_encod`` is an ``np.ndarray`` of shape ``(L, L)`` and dtype ``np.int16`` holding pairwise molecular-graph distances, as returned by :meth:`get_positional_encodings_and_align`. - ``reg`` is an ``np.ndarray`` of shape ``(n_global_regr + L - 1,)`` containing first the ``n_global_regr`` sequence-level regression targets followed by ``L - 1`` token-level targets (one per non-CLS token). Ignored when ``n_global_regr == 0``. All sequences are right-padded to the length of the longest sequence in the batch (``L_max``): - ``input_ids`` is padded with ``PAD_TOKEN_ID``. - ``attention_mask`` is padded with ``0``. - ``pos_encod`` is padded with ``np.iinfo(np.int16).max``; the diagonal of the padded region is set to ``0`` to be consistent with real token self-distances. - ``labels`` (when present) is padded with ``float("nan")`` so that padding positions can be masked out in the loss. Args: batch: List of ``(token_dict, pos_encod, reg)`` tuples, one per sample. n_global_regr: Number of sequence-level regression targets at the start of each ``reg`` array. When ``0``, no ``"labels"`` key is included in the returned dict. PAD_TOKEN_ID: Token id used to fill padded positions in ``input_ids``. Defaults to ``2``. Returns: A dict with the following keys: - ``"input_ids"`` — ``torch.LongTensor`` of shape ``(B, L_max)``. - ``"attention_mask"`` — ``torch.LongTensor`` of shape ``(B, L_max)``; ``1`` for real tokens, ``0`` for padding. - ``"positional_encodings"`` — ``torch.ShortTensor`` of shape ``(B, L_max, L_max)``. - ``"labels"`` *(only when* ``n_global_regr > 0`` *)* — ``torch.FloatTensor`` of shape ``(B, n_global_regr + L_max - 1)``; ``nan`` for padding positions. """ token_dicts, pos_encod, regs = zip(*batch) lengths = [td["input_ids"].shape[0] for td in token_dicts] L_max = max(lengths) B = len(batch) input_ids_out = np.full((B, L_max), PAD_TOKEN_ID, dtype=np.int64) attn_mask_out = np.zeros((B, L_max), dtype=np.int64) pos_encod_out = np.full((B, L_max, L_max), np.iinfo(np.int16).max, dtype=np.int16) if n_global_regr > 0: reg_out = np.full((B, n_global_regr + L_max - 1), float("nan"), dtype=np.float32) # Set diagonal to 0 up-front for the full L_max grid; individual items # already have their diagonal zeroed — this covers the padded extension. diag_idx = np.arange(L_max) pos_encod_out[:, diag_idx, diag_idx] = 0 for i, (td, pe, reg) in enumerate(zip(token_dicts, pos_encod, regs)): L = lengths[i] # Token ids & attention mask input_ids_out[i, :L] = td["input_ids"] attn_mask_out[i, :L] = td["attention_mask"] # Positional embedding (L x L) pos_encod_out[i, :L, :L] = pe # Regression: global part + token part (length L - 1, excluding CLS) if n_global_regr > 0: reg_out[i, :n_global_regr] = reg[:n_global_regr] reg_out[i, n_global_regr:n_global_regr + L - 1] = reg[n_global_regr:] out = { "input_ids": torch.from_numpy(input_ids_out), "attention_mask": torch.from_numpy(attn_mask_out), "positional_encodings": torch.from_numpy(pos_encod_out), } if n_global_regr > 0: out["labels"] = torch.from_numpy(reg_out) return out class M5EncoderModel(T5EncoderModel): def __init__(self, config: T5Config): super().__init__(config) encoder_config = config encoder_config.use_cache = False encoder_config.is_encoder_decoder = False self.encoder = M5Stack(encoder_config, self.shared) # Initialize weights and apply final processing self.post_init() def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, head_mask: Optional[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, relative_position: Optional[torch.LongTensor] = None ) -> Union[tuple[torch.FloatTensor], BaseModelOutput]: r""" input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. T5 is a model with relative position embeddings so you should be able to pad the inputs on both the right and the left. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail. To know more on how to prepare `input_ids` for pretraining take a look a [T5 Training](./t5#training). Example: ```python >>> from transformers import AutoTokenizer, T5EncoderModel >>> tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-small") >>> model = T5EncoderModel.from_pretrained("google-t5/t5-small") >>> input_ids = tokenizer( ... "Studies have been shown that owning a dog is good for you", return_tensors="pt" ... ).input_ids # Batch size 1 >>> outputs = model(input_ids=input_ids) >>> last_hidden_states = outputs.last_hidden_state ```""" return_dict = return_dict if return_dict is not None else self.config.use_return_dict encoder_outputs = self.encoder( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, relative_position=relative_position.to(dtype=torch.int32) if relative_position is not None else None ) return encoder_outputs class M5Stack(T5Stack): def __init__(self, config, embed_tokens=None): super().__init__(config, embed_tokens) self.block = nn.ModuleList( [M5Block(config, has_relative_attention_bias=bool(i == 0), layer_idx=i) for i in range(config.num_layers)] ) # Initialize weights and apply final processing self.post_init() def forward( self, input_ids=None, attention_mask=None, encoder_hidden_states=None, encoder_attention_mask=None, inputs_embeds=None, head_mask=None, cross_attn_head_mask=None, past_key_values=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, cache_position=None, relative_position=None ): # Model parallel if self.model_parallel: torch.cuda.set_device(self.first_device) self.embed_tokens = self.embed_tokens.to(self.first_device) use_cache = use_cache if use_cache is not None else self.config.use_cache 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 if input_ids is not None and inputs_embeds is not None: err_msg_prefix = "decoder_" if self.is_decoder else "" raise ValueError( f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time" ) elif input_ids is not None: input_shape = input_ids.size() input_ids = input_ids.view(-1, input_shape[-1]) elif inputs_embeds is not None: input_shape = inputs_embeds.size()[:-1] else: err_msg_prefix = "decoder_" if self.is_decoder else "" raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds") if self.gradient_checkpointing and self.training: if use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." ) use_cache = False if inputs_embeds is None: if self.embed_tokens is None: raise ValueError("You have to initialize the model with valid token embeddings") inputs_embeds = self.embed_tokens(input_ids) batch_size, seq_length = input_shape if use_cache is True: if not self.is_decoder: raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder") if self.is_decoder: if use_cache and past_key_values is None: if self.config.is_encoder_decoder: past_key_values = EncoderDecoderCache( DynamicCache(config=self.config), DynamicCache(config=self.config) ) else: past_key_values = DynamicCache(config=self.config) elif not self.is_decoder: # do not pass cache object down the line for encoder stack # it messes indexing later in decoder-stack because cache object is modified in-place past_key_values = None past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 if cache_position is None: cache_position = torch.arange( past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device ) if attention_mask is None and not is_torchdynamo_compiling(): # required mask seq length can be calculated via length of past cache mask_seq_length = past_key_values_length + seq_length attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device) if self.config.is_decoder: causal_mask = self._update_causal_mask( attention_mask, inputs_embeds, cache_position, past_key_values.self_attention_cache if isinstance(past_key_values, EncoderDecoderCache) else past_key_values, output_attentions, ) elif attention_mask is not None: causal_mask = attention_mask[:, None, None, :] causal_mask = causal_mask.to(dtype=inputs_embeds.dtype) causal_mask = (1.0 - causal_mask) * torch.finfo(inputs_embeds.dtype).min else: causal_mask = None # If a 2D or 3D attention mask is provided for the cross-attention # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] if self.is_decoder and encoder_hidden_states is not None: encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) if encoder_attention_mask is None: encoder_attention_mask = torch.ones( encoder_hidden_shape, device=inputs_embeds.device, dtype=torch.long ) encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) else: encoder_extended_attention_mask = None # Prepare head mask if needed head_mask = self.get_head_mask(head_mask, self.config.num_layers) cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers) all_hidden_states = () if output_hidden_states else None all_attentions = () if output_attentions else None all_cross_attentions = () if (output_attentions and self.is_decoder) else None position_bias = None encoder_decoder_position_bias = None hidden_states = self.dropout(inputs_embeds) for i, layer_module in enumerate(self.block): layer_head_mask = head_mask[i] cross_attn_layer_head_mask = cross_attn_head_mask[i] # Model parallel if self.model_parallel: torch.cuda.set_device(hidden_states.device) # Ensure that attention_mask is always on the same device as hidden_states if causal_mask is not None: causal_mask = causal_mask.to(hidden_states.device) if position_bias is not None: position_bias = position_bias.to(hidden_states.device) if encoder_hidden_states is not None: encoder_hidden_states = encoder_hidden_states.to(hidden_states.device) if encoder_extended_attention_mask is not None: encoder_extended_attention_mask = encoder_extended_attention_mask.to(hidden_states.device) if encoder_decoder_position_bias is not None: encoder_decoder_position_bias = encoder_decoder_position_bias.to(hidden_states.device) if layer_head_mask is not None: layer_head_mask = layer_head_mask.to(hidden_states.device) if cross_attn_layer_head_mask is not None: cross_attn_layer_head_mask = cross_attn_layer_head_mask.to(hidden_states.device) if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) layer_outputs = layer_module( hidden_states, causal_mask, position_bias, encoder_hidden_states, encoder_extended_attention_mask, encoder_decoder_position_bias, # as a positional argument for gradient checkpointing layer_head_mask=layer_head_mask, cross_attn_layer_head_mask=cross_attn_layer_head_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, return_dict=return_dict, cache_position=cache_position, relative_position=relative_position ) hidden_states = layer_outputs[0] # We share the position biases between the layers - the first layer store them # layer_outputs = hidden-states, key-valPilot phaseue-states (self-attention position bias), (self-attention weights), # (cross-attention position bias), (cross-attention weights) position_bias = layer_outputs[1] if self.is_decoder and encoder_hidden_states is not None: encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2] if output_attentions: all_attentions = all_attentions + (layer_outputs[2],) if self.is_decoder: all_cross_attentions = all_cross_attentions + (layer_outputs[4],) # Model Parallel: If it's the last layer for that device, put things on the next device if self.model_parallel: for k, v in self.device_map.items(): if i == v[-1] and "cuda:" + str(k) != self.last_device: hidden_states = hidden_states.to("cuda:" + str(k + 1)) hidden_states = self.final_layer_norm(hidden_states) hidden_states = self.dropout(hidden_states) # Add last layer if output_hidden_states: all_hidden_states = all_hidden_states + (hidden_states,) if not return_dict: return tuple( v for v in [ hidden_states, past_key_values, all_hidden_states, all_attentions, all_cross_attentions, ] if v is not None ) return BaseModelOutputWithPastAndCrossAttentions( last_hidden_state=hidden_states, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_attentions, cross_attentions=all_cross_attentions, ) class M5Block(T5Block): def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None): super().__init__(config, has_relative_attention_bias, layer_idx) self.layer = nn.ModuleList() self.layer.append( M5LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx) ) if self.is_decoder: self.layer.append(M5LayerSelfAttention(config, layer_idx=layer_idx)) self.layer.append(T5LayerFF(config)) @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") def forward( self, hidden_states, attention_mask=None, position_bias=None, encoder_hidden_states=None, encoder_attention_mask=None, encoder_decoder_position_bias=None, layer_head_mask=None, cross_attn_layer_head_mask=None, past_key_values=None, use_cache=False, output_attentions=False, return_dict=True, cache_position=None, relative_position=None, ): self_attention_outputs = self.layer[0]( hidden_states, attention_mask=attention_mask, position_bias=position_bias, layer_head_mask=layer_head_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, cache_position=cache_position, relative_position=relative_position ) hidden_states = self_attention_outputs[0] attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: clamp_value = torch.where( torch.isinf(hidden_states).any(), torch.finfo(hidden_states.dtype).max - 1000, torch.finfo(hidden_states.dtype).max, ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) do_cross_attention = self.is_decoder and encoder_hidden_states is not None if do_cross_attention: cross_attention_outputs = self.layer[1]( hidden_states, key_value_states=encoder_hidden_states, attention_mask=encoder_attention_mask, position_bias=encoder_decoder_position_bias, layer_head_mask=cross_attn_layer_head_mask, past_key_values=past_key_values, query_length=cache_position[-1] + 1, use_cache=use_cache, output_attentions=output_attentions, ) hidden_states = cross_attention_outputs[0] # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: clamp_value = torch.where( torch.isinf(hidden_states).any(), torch.finfo(hidden_states.dtype).max - 1000, torch.finfo(hidden_states.dtype).max, ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) # Keep cross-attention outputs and relative position weights attention_outputs = attention_outputs + cross_attention_outputs[1:] # Apply Feed Forward layer hidden_states = self.layer[-1](hidden_states) # clamp inf values to enable fp16 training if hidden_states.dtype == torch.float16: clamp_value = torch.where( torch.isinf(hidden_states).any(), torch.finfo(hidden_states.dtype).max - 1000, torch.finfo(hidden_states.dtype).max, ) hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) outputs = (hidden_states,) return ( outputs + attention_outputs ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights) class M5LayerSelfAttention(T5LayerSelfAttention): def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None): super().__init__(config, has_relative_attention_bias, layer_idx) self.SelfAttention = M5Attention(config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx) @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") def forward( self, hidden_states, attention_mask=None, position_bias=None, layer_head_mask=None, past_key_values=None, use_cache=False, output_attentions=False, cache_position=None, relative_position=None, ): normed_hidden_states = self.layer_norm(hidden_states) attention_output = self.SelfAttention( normed_hidden_states, mask=attention_mask, position_bias=position_bias, layer_head_mask=layer_head_mask, past_key_values=past_key_values, use_cache=use_cache, output_attentions=output_attentions, cache_position=cache_position, relative_position=relative_position ) hidden_states = hidden_states + self.dropout(attention_output[0]) outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them return outputs class M5Attention(T5Attention): """ def __init__( self, config: T5Config, has_relative_attention_bias=False, layer_idx: Optional[int] = None, ): super().__init__(config, has_relative_attention_bias, layer_idx) if self.has_relative_attention_bias: self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads) else: self.elaborate = nn.Linear() """ @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") def forward( self, hidden_states, mask=None, key_value_states=None, position_bias=None, past_key_values=None, layer_head_mask=None, query_length=None, use_cache=False, output_attentions=False, cache_position=None, relative_position=None ): """ Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states). """ # Input is (batch_size, seq_length, dim) # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder) batch_size, seq_length = hidden_states.shape[:2] # if key_value_states are provided this layer is used as a cross-attention layer for the decoder is_cross_attention = key_value_states is not None query_states = self.q(hidden_states) query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache` is_updated = False if isinstance(past_key_values, EncoderDecoderCache): is_updated = past_key_values.is_updated.get(self.layer_idx) if is_cross_attention: # after the first generated id, we can subsequently re-use all key/value_states from cache curr_past_key_value = past_key_values.cross_attention_cache else: curr_past_key_value = past_key_values.self_attention_cache else: curr_past_key_value = past_key_values current_states = key_value_states if is_cross_attention else hidden_states if is_cross_attention and past_key_values is not None and is_updated: # reuse k,v, cross_attentions key_states = curr_past_key_value.layers[self.layer_idx].keys value_states = curr_past_key_value.layers[self.layer_idx].values else: key_states = self.k(current_states) value_states = self.v(current_states) key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2) if past_key_values is not None: # save all key/value_states to cache to be re-used for fast auto-regressive generation cache_position = cache_position if not is_cross_attention else None key_states, value_states = curr_past_key_value.update( key_states, value_states, self.layer_idx, {"cache_position": cache_position} ) # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache): past_key_values.is_updated[self.layer_idx] = True # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9 scores = torch.matmul(query_states, key_states.transpose(3, 2)) if position_bias is None: key_length = key_states.shape[-2] # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past) real_seq_length = query_length if query_length is not None else cache_position[-1] + 1 if not self.has_relative_attention_bias: position_bias = torch.zeros( (1, self.n_heads, seq_length, key_length), device=scores.device, dtype=scores.dtype ) if self.gradient_checkpointing and self.training: position_bias.requires_grad = True else: position_bias = self.compute_bias( real_seq_length, key_length, device=scores.device, cache_position=cache_position, relative_position=relative_position ) position_bias = position_bias[:, :, -seq_length:, :] if mask is not None: causal_mask = mask[:, :, :, : key_states.shape[-2]] position_bias = position_bias + causal_mask if self.pruned_heads: mask = torch.ones(position_bias.shape[1]) mask[list(self.pruned_heads)] = 0 position_bias_masked = position_bias[:, mask.bool()] else: position_bias_masked = position_bias scores += position_bias_masked # (batch_size, n_heads, seq_length, key_length) attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores) attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training) # Mask heads if we want to if layer_head_mask is not None: attn_weights = attn_weights * layer_head_mask attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(batch_size, -1, self.inner_dim) attn_output = self.o(attn_output) outputs = (attn_output, position_bias) if output_attentions: outputs = outputs + (attn_weights,) return outputs @staticmethod def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128): """ Adapted from Mesh Tensorflow: https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593 Translate relative position to a bucket number for relative attention. The relative position is defined as memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for small absolute relative_position and larger buckets for larger absolute relative_positions. All relative positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket. This should allow for more graceful generalization to longer sequences than the model has been trained on Args: relative_position: an int32 Tensor bidirectional: a boolean - whether the attention is bidirectional num_buckets: an integer max_distance: an integer Returns: a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets) """ # Make all positions positive, effectively using the non-bidirectional path # However, it uses positive distances instead of negative relative_position = relative_position + 1 relative_position = torch.max(relative_position, torch.zeros_like(relative_position)) # half of the buckets are for exact increments in positions max_exact = num_buckets // 2 is_small = relative_position < max_exact num_log_buckets = num_buckets - max_exact - 1 # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance relative_position_if_large = max_exact + ( torch.log(relative_position.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - num_log_buckets) ).to(torch.long) relative_position_if_large = torch.min( relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 2) ) relative_buckets = torch.where(is_small, relative_position, relative_position_if_large) # The +1 is because we added 1 at the beginning (relative_position + 1). # This special mask is the equivalent of +inf distance and is assigned # to the last bucket. special_mask = (relative_position == np.iinfo(np.int16).max+1) relative_buckets[special_mask] = num_buckets-1 return relative_buckets def compute_bias(self, query_length, key_length, device=None, cache_position=None, relative_position=None): """Compute binned relative position bias""" if device is None: device = self.relative_attention_bias.weight.device if relative_position is None: if cache_position is None: context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] else: context_position = cache_position[:, None].to(device) memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :] relative_position = memory_position - context_position # shape (query_length, key_length) # Removing relative_position calculation breaks cache_position but it's fine since the positions are precomputed anyways relative_position_bucket = self._relative_position_bucket( relative_position, # shape (query_length, key_length) bidirectional=(not self.is_decoder), num_buckets=self.relative_attention_num_buckets, max_distance=self.relative_attention_max_distance, ) values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads) values = values.permute([0, 3, 1, 2]) # shape (batch_size, num_heads, query_length, key_length) return values # RegressionHead for tasks froms groups 0, 1, 2 and 3 # Used as regression head and classification head for pretraining class M5RegressionHead(nn.Module): def __init__(self, config: T5Config): super().__init__() self.pooler = M5Pooler(config) self.transform = nn.Linear(config.d_model, config.d_model) if config.is_gated_act: self.DenseReluDense = T5DenseGatedActDense(config) else: self.DenseReluDense = T5DenseActDense(config) self.out_proj = nn.Linear(config.d_model, config.num_labels) def forward(self, input_ids: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: pooled = self.pooler(input_ids, hidden_states) pooled = self.transform(pooled) pooled = self.DenseReluDense(pooled) output = self.out_proj(pooled) return output # TokenRegressionHead for tasks from group 4 class M5TokenRegressionHead(nn.Module): def __init__(self, config: T5Config): super().__init__() # Dimension is multiplied by 2 to account for CLS dimensional embeddings. self.transform1 = nn.Linear(config.d_model*2, config.d_model) if config.is_gated_act: self.DenseReluDense1 = T5DenseGatedActDense(config) else: self.DenseReluDense1 = T5DenseActDense(config) self.transform2 = nn.Linear(config.d_model, config.d_model) if config.is_gated_act: self.DenseReluDense2 = T5DenseGatedActDense(config) else: self.DenseReluDense2 = T5DenseActDense(config) # The output has shape (num_batches, context_length, 1) because each token has a label self.output = nn.Linear(config.d_model, 1) self.config = config def forward(self, token_hidden_states: torch.Tensor) -> torch.Tensor: # Concatenate CLS token hidden states to each token hidden state #hidden_states = torch.cat([token_hidden_states, cls_hidden_states], dim=-1) cls_hidden = token_hidden_states[:, 0, :] token_hidden = token_hidden_states[:, 1:, :] cls_repeated = cls_hidden.unsqueeze(1).expand(-1, token_hidden.size(1), -1) augmented_hidden = torch.cat([token_hidden, cls_repeated], dim=-1).contiguous() transformed = self.transform1(augmented_hidden) transformed = self.DenseReluDense1(transformed) transformed = self.transform2(transformed) transformed = self.DenseReluDense2(transformed) output = self.output(transformed) output = output.squeeze(-1) # (batch_size, num_labels) # NOTE: num_labels = seq_length return output class M5PreTrainedModel(PreTrainedModel): """ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. """ config_class = T5Config base_model_prefix = "transformer" is_parallelizable = True supports_gradient_checkpointing = True _supports_quantized_cache = False # enc-dec models don't support yet _supports_static_cache = True _supports_cache_class = True _no_split_modules = ["T5Block"] _keep_in_fp32_modules = ["wo"] @property def dummy_inputs(self): input_ids = torch.tensor(DUMMY_INPUTS) input_mask = torch.tensor(DUMMY_MASK) dummy_inputs = { "decoder_input_ids": input_ids, "input_ids": input_ids, "decoder_attention_mask": input_mask, } return dummy_inputs def _init_weights(self, module): """Initialize the weights""" factor = self.config.initializer_factor # Used for testing weights initialization if isinstance(module, T5LayerNorm): module.weight.data.fill_(factor * 1.0) elif isinstance( module, (T5Model, T5ForConditionalGeneration, T5EncoderModel, T5ForQuestionAnswering), ): # Mesh TensorFlow embeddings initialization # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624 module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0) if hasattr(module, "lm_head") and not self.config.tie_word_embeddings: module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0) if hasattr(module, "qa_outputs"): module.qa_outputs.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) module.qa_outputs.bias.data.zero_() elif isinstance(module, T5ForTokenClassification): if hasattr(module, "classifier"): module.classifier.weight.data.normal_(mean=0.0, std=factor * 1.0) module.classifier.bias.data.zero_() elif isinstance(module, T5ClassificationHead): module.dense.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.dense, "bias") and module.dense.bias is not None: module.dense.bias.data.zero_() module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: module.out_proj.bias.data.zero_() elif isinstance(module, T5DenseActDense): # Mesh TensorFlow FF initialization # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56 # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89 module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.wi, "bias") and module.wi.bias is not None: module.wi.bias.data.zero_() module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) if hasattr(module.wo, "bias") and module.wo.bias is not None: module.wo.bias.data.zero_() elif isinstance(module, T5DenseGatedActDense): module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None: module.wi_0.bias.data.zero_() module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None: module.wi_1.bias.data.zero_() module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)) if hasattr(module.wo, "bias") and module.wo.bias is not None: module.wo.bias.data.zero_() elif isinstance(module, M5RegressionHead): module.transform.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.transform, "bias") and module.transform.bias is not None: module.transform.bias.data.zero_() module.out_proj.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None: module.out_proj.bias.data.zero_() elif isinstance(module, M5TokenRegressionHead): module.transform1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model*2) ** -0.5)) module.transform1.bias.data.zero_() module.transform2.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5)) module.transform2.bias.data.zero_() module.output.weight.data.normal_(mean=0.0, std=factor * ((37.84) ** -0.5)) module.output.bias.data.zero_() elif isinstance(module, T5Attention): # Mesh TensorFlow attention initialization to avoid scaling before softmax # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136 d_model = self.config.d_model key_value_proj_dim = self.config.d_kv n_heads = self.config.num_heads module.q.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5)) module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5)) module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5)) if module.has_relative_attention_bias: module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5)) def _shift_right(self, input_ids): decoder_start_token_id = self.config.decoder_start_token_id pad_token_id = self.config.pad_token_id if decoder_start_token_id is None: raise ValueError( "self.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id. " "See T5 docs for more information." ) # shift inputs to the right if is_torch_fx_proxy(input_ids): # Item assignment is not supported natively for proxies. shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id) shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) else: shifted_input_ids = input_ids.new_zeros(input_ids.shape) shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() shifted_input_ids[..., 0] = decoder_start_token_id if pad_token_id is None: raise ValueError("self.model.config.pad_token_id has to be defined.") # replace possible -100 values in labels by `pad_token_id` shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) return shifted_input_ids class M5ModelForRegression(M5PreTrainedModel): config_class = M5EncoderConfig model_type = "m5_model" def __init__( self, config: T5Config): super().__init__(config) self.encoder: M5Encoder = M5Encoder(config) self.token_reg_head: M5TokenRegressionHead = M5TokenRegressionHead(config) self.reg_head: M5RegressionHead = M5RegressionHead(config) self.init_weights() def forward(self, input_ids, attention_mask=None, relative_position=None, **kwargs): output = self.encoder(input_ids, attention_mask, relative_position=relative_position, **kwargs) hidden_states = output.last_hidden_state tokreg_head = self.token_reg_head(hidden_states) reg_head = self.reg_head(input_ids, hidden_states) concatenated_preds = torch.cat([reg_head, tokreg_head], dim=-1) return concatenated_preds