| from typing import Optional, Tuple |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers.modeling_outputs import BaseModelOutputWithPooling |
| from transformers import PreTrainedModel |
|
|
| try: |
| from .ganda_config import UnModelConfig |
| except ImportError: |
| from ganda_config import UnModelConfig |
|
|
|
|
| class UnModel(PreTrainedModel): |
| """Transformer sentence encoder with a normalized 3072-dimensional output.""" |
|
|
| config_class = UnModelConfig |
| main_input_name = "input_ids" |
|
|
| def __init__(self, config: UnModelConfig): |
| super().__init__(config) |
|
|
| self.d_model = config.hidden_size |
| self.embed = nn.Embedding( |
| config.vocab_size, |
| config.hidden_size, |
| padding_idx=config.pad_token_id, |
| ) |
| self.pos_proj = nn.Embedding( |
| config.max_position_embeddings, |
| config.hidden_size, |
| ) |
|
|
| |
| |
| |
| self.encoder_layers = nn.TransformerEncoderLayer( |
| d_model=config.hidden_size, |
| nhead=config.num_attention_heads, |
| batch_first=True, |
| norm_first=True, |
| dim_feedforward=config.intermediate_size, |
| dropout=config.dropout, |
| activation=config.hidden_act, |
| ) |
| self.encoder = nn.TransformerEncoder( |
| encoder_layer=self.encoder_layers, |
| num_layers=config.num_hidden_layers, |
| norm=nn.RMSNorm(config.hidden_size), |
| ) |
| self.prog_head = nn.Linear(config.hidden_size, config.embedding_dim) |
|
|
| self.post_init() |
|
|
| def _init_weights(self, module): |
| |
| |
| |
| return |
|
|
| def get_input_embeddings(self): |
| return self.embed |
|
|
| def set_input_embeddings(self, value): |
| self.embed = value |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.LongTensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| inputs_embeds: Optional[torch.FloatTensor] = None, |
| output_hidden_states: Optional[bool] = None, |
| output_attentions: Optional[bool] = None, |
| return_dict: Optional[bool] = None, |
| ): |
| if input_ids is None and inputs_embeds is None: |
| raise ValueError("Pass either input_ids or inputs_embeds.") |
| if input_ids is not None and inputs_embeds is not None: |
| raise ValueError("Pass input_ids or inputs_embeds, not both.") |
|
|
| if inputs_embeds is None: |
| if input_ids.ndim != 2: |
| raise ValueError( |
| f"input_ids must have shape [batch, sequence], got {tuple(input_ids.shape)}." |
| ) |
| inputs_embeds = self.embed(input_ids) |
|
|
| batch_size, sequence_length, _ = inputs_embeds.shape |
| if sequence_length > self.config.max_position_embeddings: |
| raise ValueError( |
| f"Sequence length {sequence_length} exceeds the model limit of " |
| f"{self.config.max_position_embeddings}." |
| ) |
|
|
| if attention_mask is None: |
| attention_mask = torch.ones( |
| (batch_size, sequence_length), |
| dtype=torch.long, |
| device=inputs_embeds.device, |
| ) |
| elif attention_mask.shape != (batch_size, sequence_length): |
| raise ValueError( |
| "attention_mask must have shape " |
| f"{(batch_size, sequence_length)}, got {tuple(attention_mask.shape)}." |
| ) |
|
|
| position_ids = torch.arange(sequence_length, device=inputs_embeds.device) |
| hidden_states = inputs_embeds + self.pos_proj(position_ids).unsqueeze(0) |
| src_key_padding_mask = attention_mask.eq(0) |
|
|
| output_hidden_states = ( |
| self.config.output_hidden_states |
| if output_hidden_states is None |
| else output_hidden_states |
| ) |
| output_attentions = ( |
| self.config.output_attentions |
| if output_attentions is None |
| else output_attentions |
| ) |
| return_dict = self.config.use_return_dict if return_dict is None else return_dict |
|
|
| if output_attentions: |
| raise NotImplementedError( |
| "PyTorch TransformerEncoderLayer does not expose attention weights " |
| "through this model implementation." |
| ) |
|
|
| all_hidden_states: Optional[Tuple[torch.Tensor, ...]] = None |
| if output_hidden_states: |
| history = [hidden_states] |
| for layer in self.encoder.layers: |
| hidden_states = layer( |
| hidden_states, |
| src_key_padding_mask=src_key_padding_mask, |
| ) |
| history.append(hidden_states) |
| if self.encoder.norm is not None: |
| hidden_states = self.encoder.norm(hidden_states) |
| history[-1] = hidden_states |
| all_hidden_states = tuple(history) |
| else: |
| hidden_states = self.encoder( |
| hidden_states, |
| src_key_padding_mask=src_key_padding_mask, |
| ) |
|
|
| |
| pooling_mask = attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype) |
| summed = torch.sum(hidden_states * pooling_mask, dim=1) |
| denominator = torch.clamp(pooling_mask.sum(dim=1), min=1e-9) |
| pooled = summed / denominator |
|
|
| embeddings = self.prog_head(pooled) |
| embeddings = F.normalize(embeddings) |
|
|
| if not return_dict: |
| output = (hidden_states, embeddings) |
| if all_hidden_states is not None: |
| output += (all_hidden_states,) |
| return output |
|
|
| return BaseModelOutputWithPooling( |
| last_hidden_state=hidden_states, |
| pooler_output=embeddings, |
| hidden_states=all_hidden_states, |
| ) |
|
|
| @torch.no_grad() |
| def generate_embeddings( |
| self, |
| input_ids: torch.LongTensor, |
| attention_mask: Optional[torch.Tensor] = None, |
| ) -> torch.Tensor: |
| """Return L2-normalized sentence embeddings with shape [batch, embedding_dim].""" |
| if input_ids.ndim == 1: |
| input_ids = input_ids.unsqueeze(0) |
| if attention_mask is not None and attention_mask.ndim == 1: |
| attention_mask = attention_mask.unsqueeze(0) |
| return self( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| return_dict=True, |
| ).pooler_output |
|
|
|
|
| UnModel.register_for_auto_class("AutoModel") |
|
|
