File size: 6,970 Bytes
8e1984a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | 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,
)
# Keep the encoder_layers attribute as well as encoder.layers. The
# original training class created both, so this preserves its state_dict
# key layout (including encoder_layers.*, which is not used in forward).
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):
# PyTorch layers initialize themselves on construction. Leaving those
# values untouched also preserves the original model's initialization
# behavior when a new model is instantiated before loading a checkpoint.
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,
)
# Masked mean pooling, matching the original training implementation.
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")
|