contra-bottleneck-t5-large-wikipedia / bottleneck_t5.py

Upload BottleneckT5LMWithPerturb

6986db8 8 months ago

No virus

18.9 kB

	import copy
	import torch
	import torch.nn as nn
	from torch.nn import CrossEntropyLoss
	from torch.nn import functional as F
	from typing import Optional, Tuple, Union

	from transformers import T5Config, T5Tokenizer, T5ForConditionalGeneration
	from transformers.models.t5.modeling_t5 import (
	T5LayerNorm,
	T5LayerFF,
	T5Attention,
	T5LayerSelfAttention,
	T5LayerCrossAttention,
	T5Block,
	T5Stack,
	)
	from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput

	class BottleneckCrossAttentionGate(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.gate = nn.Linear(2 * config.d_model, config.d_model, bias=False)
	self.act = nn.Sigmoid()

	def forward(self, query_states, latents):
	latents = latents.unsqueeze(1).expand(query_states.shape)
	query_latents = torch.cat([query_states, latents], dim=-1)
	return 2 * self.act(self.gate(query_latents))

	class BottleneckT5Attention(T5Attention):
	def __init__(self, config: T5Config, has_relative_attention_bias=False):
	super(T5Attention, self).__init__()
	self.is_decoder = config.is_decoder
	self.has_relative_attention_bias = has_relative_attention_bias
	self.relative_attention_num_buckets = config.relative_attention_num_buckets
	self.relative_attention_max_distance = config.relative_attention_max_distance
	self.d_model = config.d_model
	self.key_value_proj_dim = config.d_kv
	self.n_heads = config.num_heads
	self.dropout = config.dropout_rate
	self.inner_dim = self.n_heads * self.key_value_proj_dim

	# Mesh TensorFlow initialization to avoid scaling before softmax
	self.v = nn.Linear(self.d_model, self.inner_dim, bias=False)
	self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)

	if self.has_relative_attention_bias:
	self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)
	self.pruned_heads = set()
	self.gradient_checkpointing = False

	def prune_heads(self, heads):
	if len(heads) == 0:
	return
	heads, index = find_pruneable_heads_and_indices(
	heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads
	)
	# Prune linear layers
	self.v = prune_linear_layer(self.v, index)
	self.o = prune_linear_layer(self.o, index, dim=1)
	# Update hyper params
	self.n_heads = self.n_heads - len(heads)
	self.inner_dim = self.key_value_proj_dim * self.n_heads
	self.pruned_heads = self.pruned_heads.union(heads)

	def forward(
	self,
	hidden_states,
	mask=None,
	key_value_states=None,
	position_bias=None,
	past_key_value=None,
	layer_head_mask=None,
	query_length=None,
	use_cache=False,
	output_attentions=False,
	):
	"""
	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, key_length) (non-causal) or (batch_size, key_length, key_length)
	# past_key_value[0] is (batch_size, n_heads, q_len - 1, dim_per_head)
	batch_size, seq_length = hidden_states.shape[:2]

	real_seq_length = seq_length

	if past_key_value is not None:
	assert (
	len(past_key_value) == 2
	), f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
	real_seq_length += past_key_value[0].shape[2] if query_length is None else query_length

	key_length = real_seq_length if key_value_states is None else key_value_states.shape[1]

	def shape(states):
	"""projection"""
	return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)

	def unshape(states):
	"""reshape"""
	return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)

	def project(hidden_states, proj_layer, key_value_states, past_key_value):
	"""projects hidden states correctly to key/query states"""
	if key_value_states is None:
	# self-attn
	# (batch_size, n_heads, seq_length, dim_per_head)
	hidden_states = shape(proj_layer(hidden_states))
	elif past_key_value is None:
	# cross-attn
	# (batch_size, n_heads, seq_length, dim_per_head)
	hidden_states = shape(proj_layer(key_value_states))

	if past_key_value is not None:
	if key_value_states is None:
	# self-attn
	# (batch_size, n_heads, key_length, dim_per_head)
	hidden_states = torch.cat([past_key_value, hidden_states], dim=2)
	else:
	# cross-attn
	hidden_states = past_key_value
	return hidden_states

	# key/value states
	key_states = torch.zeros((batch_size, self.n_heads, seq_length, key_length), device=hidden_states.device)
	value_states = project(
	hidden_states, self.v, key_value_states, past_key_value[1] if past_key_value is not None else None
	)

	# compute scores
	scores = torch.ones((batch_size, self.n_heads, seq_length, key_length), device=hidden_states.device)

	if position_bias is None:
	if not self.has_relative_attention_bias:
	position_bias = torch.zeros(
	(1, self.n_heads, real_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)

	# if key and values are already calculated
	# we want only the last query position bias
	if past_key_value is not None:
	position_bias = position_bias[:, :, -hidden_states.size(1) :, :]

	if mask is not None:
	position_bias = position_bias + mask # (batch_size, n_heads, seq_length, key_length)

	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
	attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
	scores
	) # (batch_size, n_heads, seq_length, key_length)
	attn_weights = nn.functional.dropout(
	attn_weights, p=self.dropout, training=self.training
	) # (batch_size, n_heads, seq_length, key_length)

	# Mask heads if we want to
	if layer_head_mask is not None:
	attn_weights = attn_weights * layer_head_mask

	attn_output = unshape(torch.matmul(attn_weights, value_states)) # (batch_size, seq_length, dim)
	attn_output = self.o(attn_output)

	present_key_value_state = (key_states, value_states) if (self.is_decoder and use_cache) else None
	outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)

	if output_attentions:
	outputs = outputs + (attn_weights,)
	return outputs

	class BottleneckT5LayerCrossAttention(T5LayerCrossAttention):
	def __init__(self, config):
	super(T5LayerCrossAttention, self).__init__()
	self.EncDecAttention = BottleneckT5Attention(config, has_relative_attention_bias=False)
	self.layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
	self.gate = BottleneckCrossAttentionGate(config)
	self.dropout = nn.Dropout(config.dropout_rate)

	def forward(
	self,
	hidden_states,
	key_value_states,
	attention_mask=None,
	position_bias=None,
	layer_head_mask=None,
	past_key_value=None,
	use_cache=False,
	query_length=None,
	output_attentions=False,
	):
	normed_hidden_states = self.layer_norm(hidden_states)
	attention_output = self.EncDecAttention(
	normed_hidden_states,
	mask=attention_mask,
	key_value_states=key_value_states,
	position_bias=position_bias,
	layer_head_mask=layer_head_mask,
	past_key_value=past_key_value,
	use_cache=use_cache,
	query_length=query_length,
	output_attentions=output_attentions,
	)
	latents = key_value_states[:, 0]
	layer_output = hidden_states + self.dropout(self.gate(normed_hidden_states, latents) * attention_output[0])
	outputs = (layer_output,) + attention_output[1:] # add attentions if we output them
	return outputs

	class BottleneckT5Block(T5Block):
	def __init__(self, config, has_relative_attention_bias=False):
	super(T5Block, self).__init__()
	self.is_decoder = config.is_decoder
	self.layer = nn.ModuleList()
	self.layer.append(T5LayerSelfAttention(config, has_relative_attention_bias=has_relative_attention_bias))
	if self.is_decoder:
	self.layer.append(BottleneckT5LayerCrossAttention(config))

	self.layer.append(T5LayerFF(config))

	class BottleneckT5Stack(T5Stack):
	def __init__(self, config, embed_tokens=None):
	super(T5Stack, self).__init__(config)

	self.embed_tokens = embed_tokens
	self.is_decoder = config.is_decoder

	self.block = nn.ModuleList(
	[BottleneckT5Block(config, has_relative_attention_bias=bool(i == 0)) for i in range(config.num_layers)]
	)
	self.final_layer_norm = T5LayerNorm(config.d_model, eps=config.layer_norm_epsilon)
	self.dropout = nn.Dropout(config.dropout_rate)

	# Initialize weights and apply final processing
	self.post_init()
	# Model parallel
	self.model_parallel = False
	self.device_map = None
	self.gradient_checkpointing = False

	class BottleneckT5LMWithPerturb(T5ForConditionalGeneration):
	def __init__(self, config: T5Config):
	super(T5ForConditionalGeneration, self).__init__(config)
	self.model_dim = config.d_model

	self.shared = nn.Embedding(config.vocab_size, config.d_model)
	encoder_config = copy.deepcopy(config)
	encoder_config.is_decoder = False
	encoder_config.use_cache = False
	encoder_config.is_encoder_decoder = False
	self.encoder = T5Stack(encoder_config, self.shared)

	# New in Contra: MHA bottleneck block
	self.num_heads = config.num_heads
	self.bottleneck = nn.MultiheadAttention(config.d_model,
	num_heads=config.num_heads,
	dropout=config.dropout_rate,
	bias=False,
	batch_first=True)
	self.bottleneck_scale = nn.Parameter(torch.ones(1))

	self.dec_emb = nn.Embedding(config.vocab_size, config.d_model)
	decoder_config = copy.deepcopy(config)
	decoder_config.is_decoder = True
	decoder_config.is_encoder_decoder = False
	decoder_config.num_layers = config.num_decoder_layers
	self.decoder = BottleneckT5Stack(decoder_config, self.dec_emb)

	self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

	# Initialize weights and apply final processing
	self.post_init()

	# Model parallel
	self.model_parallel = False
	self.device_map = None

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	attention_mask: Optional[torch.FloatTensor] = None,
	decoder_input_ids: Optional[torch.LongTensor] = None,
	decoder_attention_mask: Optional[torch.BoolTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	decoder_head_mask: Optional[torch.FloatTensor] = None,
	cross_attn_head_mask: Optional[torch.Tensor] = None,
	encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None,
	past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
	labels: Optional[torch.LongTensor] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	perturb_vector: Optional[torch.FloatTensor] = None,
	encode_only: Optional[bool] = None,
	) -> Union[Tuple[torch.FloatTensor], Seq2SeqLMOutput]:
	use_cache = use_cache if use_cache is not None else self.config.use_cache
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	# FutureWarning: head_mask was separated into two input args - head_mask, decoder_head_mask
	if head_mask is not None and decoder_head_mask is None:
	if self.config.num_layers == self.config.num_decoder_layers:
	warnings.warn(__HEAD_MASK_WARNING_MSG, FutureWarning)
	decoder_head_mask = head_mask

	# Encode if needed (training, first prediction pass)
	if encoder_outputs is None:
	# Convert encoder inputs in embeddings if needed
	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,
	)
	elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):
	encoder_outputs = BaseModelOutput(
	last_hidden_state=encoder_outputs[0],
	hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
	attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
	)

	hidden_states = encoder_outputs[0]

	# MHA across token embeddings + embedding normalization + broadcast
	hidden_states = hidden_states.repeat(
	attention_mask.shape[0] // hidden_states.shape[0],
	1, 1) # during contrastive search, attn mask can have higher batch size than hidden_state
	mask_expanded = attention_mask.float().unsqueeze(-1).expand(hidden_states.shape)
	mean_pooled_embedding = torch.sum(hidden_states * mask_expanded, 1) / torch.clamp(mask_expanded.sum(1), min=1e-9)
	unscaled_latent, attn_weights = self.bottleneck(mean_pooled_embedding.unsqueeze(1), hidden_states, hidden_states,
	need_weights=False,
	# torch MHA attn_mask has opposite signs to HF T5 masks... sigh
	attn_mask=attention_mask.float().unsqueeze(1).repeat_interleave(self.num_heads, dim=0))
	latent = self.bottleneck_scale * F.normalize(unscaled_latent, p=2, dim=2)
	if encode_only:
	return latent.squeeze(1)
	hidden_states = latent.expand(hidden_states.shape)

	if hasattr(self, 'perturb_vector'):
	hidden_states = self.bottleneck_scale * F.normalize(hidden_states + self.perturb_vector, p=2, dim=2)

	if self.model_parallel:
	torch.cuda.set_device(self.decoder.first_device)

	if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
	# get decoder inputs from shifting lm labels to the right
	decoder_input_ids = self._shift_right(labels)

	# Set device for model parallelism
	if self.model_parallel:
	torch.cuda.set_device(self.decoder.first_device)
	hidden_states = hidden_states.to(self.decoder.first_device)
	if decoder_input_ids is not None:
	decoder_input_ids = decoder_input_ids.to(self.decoder.first_device)
	if attention_mask is not None:
	attention_mask = attention_mask.to(self.decoder.first_device)
	if decoder_attention_mask is not None:
	decoder_attention_mask = decoder_attention_mask.to(self.decoder.first_device)

	# Decode
	decoder_outputs = self.decoder(
	input_ids=decoder_input_ids,
	attention_mask=decoder_attention_mask,
	inputs_embeds=decoder_inputs_embeds,
	past_key_values=past_key_values,
	encoder_hidden_states=hidden_states,
	encoder_attention_mask=attention_mask,
	head_mask=decoder_head_mask,
	cross_attn_head_mask=cross_attn_head_mask,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	sequence_output = decoder_outputs[0]

	# Set device for model parallelism
	if self.model_parallel:
	torch.cuda.set_device(self.encoder.first_device)
	self.lm_head = self.lm_head.to(self.encoder.first_device)
	sequence_output = sequence_output.to(self.lm_head.weight.device)

	if self.config.tie_word_embeddings:
	# Rescale output before projecting on vocab
	# See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
	sequence_output = sequence_output * (self.model_dim**-0.5)

	lm_logits = self.lm_head(sequence_output)

	loss = None
	if labels is not None:
	loss_fct = CrossEntropyLoss(ignore_index=-100)
	loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
	# TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666

	if not return_dict:
	output = (lm_logits,) + decoder_outputs[1:] + encoder_outputs
	return ((loss,) + output) if loss is not None else output

	return Seq2SeqLMOutput(
	loss=loss,
	logits=lm_logits,
	past_key_values=decoder_outputs.past_key_values,
	decoder_hidden_states=decoder_outputs.hidden_states,
	decoder_attentions=decoder_outputs.attentions,
	cross_attentions=decoder_outputs.cross_attentions,
	encoder_last_hidden_state=encoder_outputs.last_hidden_state,
	encoder_hidden_states=encoder_outputs.hidden_states,
	encoder_attentions=encoder_outputs.attentions,
	)