lsg-legal-small-uncased-4096 / modeling_lsg_bert.py

small fix

b6adbe7 5 months ago

No virus

47.2 kB

	from logging import warn
	from transformers.models.bert.modeling_bert import *
	import torch
	import torch.nn as nn
	from transformers.models.bert.configuration_bert import BertConfig
	import sys

	AUTO_MAP = {
	"AutoModel": "modeling_lsg_bert.LSGBertModel",
	"AutoModelForCausalLM": "modeling_lsg_bert.LSGBertLMHeadModel",
	"AutoModelForMaskedLM": "modeling_lsg_bert.LSGBertForMaskedLM",
	"AutoModelForPreTraining": "modeling_lsg_bert.LSGBertForPreTraining",
	"AutoModelForMultipleChoice": "modeling_lsg_bert.LSGBertForMultipleChoice",
	"AutoModelForQuestionAnswering": "modeling_lsg_bert.LSGBertForQuestionAnswering",
	"AutoModelForSequenceClassification": "modeling_lsg_bert.LSGBertForSequenceClassification",
	"AutoModelForTokenClassification": "modeling_lsg_bert.LSGBertForTokenClassification"
	}

	class LSGBertConfig(BertConfig):
	"""
	This class overrides :class:`~transformers.BertConfig`. Please check the superclass for the appropriate
	documentation alongside usage examples.
	"""

	base_model_prefix = "lsg"
	model_type = "bert"

	def __init__(
	self,
	adaptive=True,
	base_model_prefix="lsg",
	block_size=128,
	lsh_num_pre_rounds=1,
	mask_first_token=False,
	num_global_tokens=1,
	pool_with_global=True,
	sparse_block_size=128,
	sparsity_factor=2,
	sparsity_type="norm",
	**kwargs
	):
	"""Constructs LSGBertConfig."""
	super().__init__(**kwargs)

	self.adaptive = adaptive
	self.auto_map = AUTO_MAP
	self.base_model_prefix = base_model_prefix
	self.block_size = block_size
	self.lsh_num_pre_rounds = lsh_num_pre_rounds
	self.mask_first_token = mask_first_token
	self.num_global_tokens = num_global_tokens
	self.pool_with_global = pool_with_global
	self.sparse_block_size = sparse_block_size
	self.sparsity_factor = sparsity_factor
	self.sparsity_type = sparsity_type

	if sparsity_type not in [None, "none", "norm", "lsh", "pooling", "stride", "block_stride", "bos_pooling"]:
	logger.warning(
	"[WARNING CONFIG]: sparsity_mode not in [None, 'none', 'norm', 'lsh', 'pooling', 'stride', 'block_stride', 'bos_pooling'], \
	setting sparsity_type=None, computation will skip sparse attention")
	self.sparsity_type = None

	if self.sparsity_type in ["stride", "block_stride"]:
	if self.sparsity_factor > self.num_attention_heads:
	logger.warning(
	"[WARNING CONFIG]: sparsity_factor > num_attention_heads is not recommended for stride/block_stride sparsity"
	)

	if self.num_global_tokens < 1:
	logger.warning(
	"[WARNING CONFIG]: num_global_tokens < 1 is not compatible, setting num_global_tokens=1"
	)
	self.num_global_tokens = 1
	elif self.num_global_tokens > 512:
	logger.warning(
	"[WARNING CONFIG]: num_global_tokens > 512 is not allowed, setting num_global_tokens=512"
	)
	self.num_global_tokens = 512

	if self.sparsity_factor > 0:
	assert self.block_size % self.sparsity_factor == 0, "[ERROR CONFIG]: block_size must be divisible by sparsity_factor"
	assert self.block_size//self.sparsity_factor >= 1, "[ERROR CONFIG]: make sure block_size >= sparsity_factor"

	if self.mask_first_token and not pool_with_global:
	logger.warning(
	"[WARNING CONFIG]: pool_with_global==False is not compatible with mask_first_token==True. Setting pool_with_global to True.")
	self.pool_with_global = True

	if hasattr(self, "position_embedding_type"):
	if self.position_embedding_type != "absolute":
	logger.warning(
	"[WARNING CONFIG]: LSG Attention is not compatible with relative positional embedding and will skip its computation. Set position_embedding_type='absolute' to remove this warning.")


	class BaseSelfAttention(nn.Module):

	def init_modules(self, config):
	if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
	config, "embedding_size"
	):
	raise ValueError(
	"The hidden size (%d) is not a multiple of the number of attention "
	"heads (%d)" % (config.hidden_size, config.num_attention_heads)
	)

	self.num_attention_heads = config.num_attention_heads
	self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
	self.all_head_size = self.num_attention_heads * self.attention_head_size

	self.query = nn.Linear(config.hidden_size, self.all_head_size)
	self.key = nn.Linear(config.hidden_size, self.all_head_size)
	self.value = nn.Linear(config.hidden_size, self.all_head_size)

	self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

	def transpose_for_scores(self, x):
	new_x_shape = x.size()[:-1] + (
	self.num_attention_heads,
	self.attention_head_size,
	)
	x = x.view(*new_x_shape)
	return x.permute(0, 2, 1, 3)

	def reshape_output(self, context_layer):
	context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
	new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
	return context_layer.view(*new_context_layer_shape)

	def project_QKV(self, hidden_states):

	query_layer = self.transpose_for_scores(self.query(hidden_states))
	key_layer = self.transpose_for_scores(self.key(hidden_states))
	value_layer = self.transpose_for_scores(self.value(hidden_states))
	return query_layer, key_layer, value_layer


	class BaseAttentionProduct(nn.Module):

	def __init__(self, config):
	"""
	Compute attention: softmax(Q @ K.T) @ V
	"""
	super().__init__()
	self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

	def forward(self, query_layer, key_layer, value_layer, attention_mask=None):

	d = query_layer.shape[-1]

	# Take the dot product between "query" and "key" to get the raw attention scores.
	attention_scores = query_layer @ key_layer.transpose(-1, -2) / math.sqrt(d)

	del query_layer
	del key_layer

	if attention_mask is not None:
	# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
	attention_scores = attention_scores + attention_mask
	del attention_mask

	# Normalize the attention scores to probabilities.
	attention_probs = nn.Softmax(dim=-1)(attention_scores)

	# This is actually dropping out entire tokens to attend to, which might
	# seem a bit unusual, but is taken from the original Transformer paper.
	context_layer = self.dropout(attention_probs) @ value_layer

	return context_layer


	class CausalAttentionProduct(nn.Module):

	def __init__(self, config):
	"""
	Compute attention: softmax(Q @ K.T) @ V
	"""
	super().__init__()
	self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
	self.block_size = config.block_size

	def forward(self, query_layer, key_layer, value_layer, attention_mask=None, causal_shape=None):

	d = query_layer.shape[-1]

	# Take the dot product between "query" and "key" to get the raw attention scores.
	attention_scores = query_layer @ key_layer.transpose(-1, -2) / math.sqrt(d)

	del query_layer
	del key_layer

	if attention_mask is not None:
	# Add causal mask
	causal_shape = (self.block_size, self.block_size) if causal_shape is None else causal_shape
	causal_mask = torch.tril(
	torch.ones(*causal_shape, device=attention_mask.device, dtype=attention_scores.dtype),
	diagonal=-1
	)

	# Min value
	dtype_min = torch.tensor(
	torch.finfo(attention_scores.dtype).min, device=attention_scores.device, dtype=attention_scores.dtype
	)

	# Build causal + attention_mask
	causal_mask = torch.nn.functional.pad(causal_mask.T * dtype_min, (attention_mask.size()[-1] - self.block_size, 0), value=0)
	attention_mask = torch.max(attention_mask + causal_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0), dtype_min)

	attention_scores = attention_scores + attention_mask
	del attention_mask
	del causal_mask

	# Normalize the attention scores to probabilities.
	attention_probs = nn.Softmax(dim=-1)(attention_scores)

	# This is actually dropping out entire tokens to attend to, which might
	# seem a bit unusual, but is taken from the original Transformer paper.
	context_layer = self.dropout(attention_probs) @ value_layer

	return context_layer


	class LSGAttentionProduct(nn.Module):

	def __init__(self, config, block_size=None, sparse_block_size=None, sparsity_factor=4, is_causal=False):
	"""
	Compute block or overlapping blocks attention products
	"""
	super().__init__()

	self.block_size = block_size
	self.sparse_block_size = sparse_block_size
	self.sparsity_factor = sparsity_factor
	self.is_causal = is_causal

	if self.block_size is None:
	self.block_size = config.block_size

	if self.sparse_block_size is None:
	self.sparse_block_size = config.sparse_block_size

	# Shape of blocks
	self.local_shapes = (self.block_size*3, self.block_size)
	if self.sparse_block_size and self.sparsity_factor > 0:
	self.sparse_shapes = (self.sparse_block_size*3, self.block_size//self.sparsity_factor)

	if is_causal:
	self.attention = CausalAttentionProduct(config)
	else:
	self.attention = BaseAttentionProduct(config)

	def build_lsg_inputs(self, hidden_states, sparse_hidden_states, global_hidden_states, is_attn_mask=False):

	# Build local tokens
	local_hidden_states = self.reshape_to_local_block(hidden_states, is_attn_mask)
	del hidden_states

	# Build sparse tokens
	if sparse_hidden_states is not None:
	sparse_hidden_states = self.reshape_to_sparse_block(sparse_hidden_states, is_attn_mask)

	return self.cat_global_sparse_local_tokens(global_hidden_states, sparse_hidden_states, local_hidden_states)

	def forward(
	self,
	query_layer,
	key_layer,
	value_layer,
	attention_mask=None,
	sparse_key=None,
	sparse_value=None,
	sparse_mask=None,
	global_key=None,
	global_value=None,
	global_mask=None
	):

	# Input batch, heads, length, hidden_size
	n, h, t, d = query_layer.size()
	n_blocks = t // self.block_size
	assert t % self.block_size == 0

	key_layer = self.build_lsg_inputs(
	key_layer,
	sparse_key,
	global_key
	)
	del sparse_key
	del global_key

	value_layer = self.build_lsg_inputs(
	value_layer,
	sparse_value,
	global_value
	)
	del sparse_value
	del global_value

	attention_mask = self.build_lsg_inputs(
	attention_mask,
	sparse_mask,
	global_mask.transpose(-1, -2),
	is_attn_mask=True
	).transpose(-1, -2)
	del sparse_mask
	del global_mask

	# expect (..., t, d) shape
	# Compute attention
	context_layer = self.attention(
	query_layer=self.chunk(query_layer, n_blocks),
	key_layer=key_layer,
	value_layer=value_layer,
	attention_mask=attention_mask
	)

	return context_layer.reshape(n, h, -1, d)

	def reshape_to_local_block(self, hidden_states, is_attn_mask=False):

	size, step = self.local_shapes
	s = (size - step) // 2

	# Pad before block reshaping
	if is_attn_mask:
	pad_value = torch.finfo(hidden_states.dtype).min
	hidden_states = hidden_states.transpose(-1, -2)
	else:
	pad_value = 0

	hidden_states = torch.nn.functional.pad(
	hidden_states.transpose(-1, -2),
	pad=(s, s),
	value=pad_value
	).transpose(-1, -2)

	# Make blocks
	hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)

	# Skip third block if causal
	if self.is_causal:
	return hidden_states[..., :size*2//3, :]

	return hidden_states

	def reshape_to_sparse_block(self, hidden_states, is_attn_mask=False):

	size, step = self.sparse_shapes

	# In case of odd case
	odd_offset = (step % 2)

	# n, h, t, d*2 + 1
	size = size*2
	s = (size - step) // 2 + odd_offset

	# Pad before block reshaping
	if is_attn_mask:
	pad_value = torch.finfo(hidden_states.dtype).min
	hidden_states = hidden_states.transpose(-1, -2)
	else:
	pad_value = 0

	hidden_states = torch.nn.functional.pad(
	hidden_states.transpose(-1, -2),
	pad=(s, s),
	value=pad_value
	).transpose(-1, -2)

	# Make blocks
	hidden_states = hidden_states.unfold(-2, size=size, step=step).transpose(-1, -2)

	# Fix case where block_size == sparsify_factor
	if odd_offset:
	hidden_states = hidden_states[..., :-1, :, :]

	# Indexes for selection
	u = (size - self.block_size * 3 // self.sparsity_factor) // 2 + odd_offset
	s = self.sparse_block_size

	# Skip right block if causal
	if self.is_causal:
	return hidden_states[..., u-s:u, :]

	u_ = u + odd_offset
	return torch.cat([hidden_states[..., u-s:u, :], hidden_states[..., -u_:-u_+s, :]], dim=-2)

	def cat_global_sparse_local_tokens(self, x_global, x_sparse=None, x_local=None, dim=-2):

	n, h, b, t, d = x_local.size()
	x_global = x_global.unsqueeze(-3).expand(-1, -1, b, -1, -1)
	if x_sparse is not None:
	return torch.cat([x_global, x_sparse, x_local], dim=dim)
	return torch.cat([x_global, x_local], dim=dim)

	def chunk(self, x, n_blocks):

	t, d = x.size()[-2:]
	return x.reshape(*x.size()[:-2], n_blocks, -1, d)


	class LSGBertEmbeddings(BertEmbeddings):

	def __init__(self, config):
	super().__init__(config)

	self.num_global_tokens = config.num_global_tokens

	# Hardcoded but partially trained
	self.global_embeddings = nn.Embedding(512, embedding_dim=config.hidden_size, )

	self.block_size = config.block_size

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	token_type_ids: Optional[torch.LongTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	past_key_values_length: int = 0,
	) -> torch.Tensor:
	if input_ids is not None:
	input_shape = input_ids.size()
	else:
	input_shape = inputs_embeds.size()[:-1]

	seq_length = input_shape[1]

	if position_ids is None:
	position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]

	# Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
	# when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
	# issue #5664
	if token_type_ids is None:
	if hasattr(self, "token_type_ids"):
	buffered_token_type_ids = self.token_type_ids[:, :seq_length]
	buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
	token_type_ids = buffered_token_type_ids_expanded
	else:
	token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

	if inputs_embeds is None:
	inputs_embeds = self.word_embeddings(input_ids)
	token_type_embeddings = self.token_type_embeddings(token_type_ids[:, :seq_length])

	embeddings = inputs_embeds + token_type_embeddings
	if self.position_embedding_type == "absolute":
	position_embeddings = self.position_embeddings(position_ids[:, :seq_length])
	embeddings += position_embeddings

	#if self.num_global_tokens < 0:
	n, t, d = embeddings.size()

	# Add global_tokens
	indexes = torch.arange(self.num_global_tokens, device=embeddings.device).reshape(1, -1)
	global_embeddings = self.global_embeddings(indexes)
	embeddings = torch.cat([global_embeddings.expand(n, -1, d), embeddings], dim=-2)

	embeddings = self.LayerNorm(embeddings)
	embeddings = self.dropout(embeddings)
	return embeddings


	class LSGSelfAttention(BaseSelfAttention):
	'''
	Compute local attention with overlapping blocs
	Use global attention for tokens with highest norm
	'''
	def __init__(self, config):
	super().__init__()

	self.init_modules(config)

	self.block_size = config.block_size
	self.sparse_block_size = config.sparse_block_size
	self.num_global_tokens = config.num_global_tokens
	self.sparsity_factor = config.sparsity_factor
	self.is_causal = config.is_decoder
	self.is_decoder = config.is_decoder

	self.attention = LSGAttentionProduct(
	config,
	block_size=config.block_size,
	sparse_block_size=config.sparse_block_size,
	sparsity_factor=self.sparsity_factor,
	is_causal=self.is_causal
	)

	if self.is_causal:
	self.causal_attention = CausalAttentionProduct(config)
	self.full_attention = BaseAttentionProduct(config)

	sparse_functions = {
	"norm": self.get_sparse_tokens_with_norm,
	"pooling": self.get_sparse_tokens_with_pooling,
	"lsh": self.get_sparse_tokens_with_lsh,
	"stride": self.get_sparse_tokens_with_stride,
	"block_stride": self.get_sparse_tokens_with_block_stride,
	"bos_pooling": self.get_sparse_tokens_with_bos_pooling
	}

	self.sparsity_type = config.sparsity_type
	self.get_sparse_elements = sparse_functions.get(self.sparsity_type, lambda w, x, y, z: (None, None, None))

	if config.sparsity_type == "lsh":
	self.lsh_num_pre_rounds = config.lsh_num_pre_rounds

	def get_sparse_tokens_with_norm(self, queries, keys, values, mask):

	if self.sparsity_factor == 1:
	return keys, values, mask.expand(-1, keys.size()[1], -1, -1)

	with torch.no_grad():

	block_size = min(self.block_size, self.sparse_block_size)
	key_norm = keys.detach().norm(dim=-1, keepdim=True)
	key_norm = key_norm * ~mask.transpose(-1, -2).bool()
	key_norm = self.chunk(key_norm, block_size)

	n, h, b, t, d = key_norm.size()

	idx = key_norm.argsort(dim=-2)
	del key_norm
	idx += (torch.arange(b, device=keys.device)*t).reshape(1, 1, b, 1, 1)

	split = (t - block_size // self.sparsity_factor, block_size // self.sparsity_factor)
	sparse_idx = idx.split(split, -2)[-1].reshape(n, h, -1, 1)

	d = keys.size()[-1]
	keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
	values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
	mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)

	return keys, values, mask

	def get_sparse_tokens_with_pooling(self, queries, keys, values, mask):

	if self.sparsity_factor == 1:
	return keys, values, mask.expand(-1, keys.size()[1], -1, -1)

	keys = self.chunk(keys, self.sparsity_factor)
	values = self.chunk(values, self.sparsity_factor)

	n, h, b, t, d = keys.size()
	mask = mask.reshape(n, 1, b, 1, t)
	mask = ~mask.transpose(-1, -2).bool()

	keys = keys * mask
	values = values * mask

	mask = mask.sum(dim=-2)
	keys = keys.sum(dim=-2) / (mask + 1e-6)
	values = values.sum(dim=-2) / (mask + 1e-6)

	mask = (1. - mask.clamp(0, 1))
	mask *= torch.finfo(mask.dtype).min
	return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)

	def get_sparse_tokens_with_stride(self, queries, keys, values, mask):

	if self.sparsity_factor == 1:
	return keys, values, mask.expand(-1, keys.size()[1], -1, -1)

	n, h, t, d = keys.size()
	sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device) * self.sparsity_factor
	sparse_idx = sparse_idx.reshape(1, 1, -1, 1) + (torch.arange(h, device=keys.device) % self.sparsity_factor).reshape(1, h, 1, 1)
	sparse_idx = sparse_idx.expand(n, h, -1, 1)

	keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
	values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
	mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)

	return keys, values, mask

	def get_sparse_tokens_with_block_stride(self, queries, keys, values, mask):

	if self.sparsity_factor == 1:
	return keys, values, mask.expand(-1, keys.size()[1], -1, -1)

	n, h, t, d = keys.size()

	t, b = self.block_size, t // self.block_size
	sparse_idx = torch.arange(t // self.sparsity_factor, device=keys.device)
	sparse_idx = sparse_idx.reshape(1, 1, 1, -1, 1) + torch.arange(h, device=keys.device).reshape(1, h, 1, 1, 1) * (t // self.sparsity_factor)
	sparse_idx = (sparse_idx % t)
	sparse_idx = sparse_idx + torch.arange(b, device=keys.device).reshape(1, 1, -1, 1, 1) * t
	sparse_idx = sparse_idx.reshape(1, h, -1, 1).expand(n, h, -1, 1)

	keys = keys.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
	values = values.gather(dim=-2, index=sparse_idx.expand(-1, -1, -1, d))
	mask = mask.expand(-1, h, -1, -1).transpose(-1, -2).gather(dim=-2, index=sparse_idx).transpose(-1, -2)

	return keys, values, mask

	def get_sparse_tokens_with_lsh(self, queries, keys, values, mask):

	if self.sparsity_factor == 1:
	return keys, values, mask.expand(-1, keys.size()[1], -1, -1)

	if self.sparsity_factor == self.sparse_block_size:
	return self.get_sparse_tokens_with_bos_pooling(queries, keys, values, mask)

	block_size = min(self.block_size, self.sparse_block_size)
	keys = self.chunk(keys, block_size)
	values = self.chunk(values, block_size)

	n, h, b, t, d = keys.size()
	mask = mask.reshape(n, 1, b, 1, t)
	mask = ~mask.transpose(-1, -2).bool()

	keys = keys * mask
	values = values * mask
	mask = mask.expand(-1, h, -1, -1, -1).float()

	extra_factor = 1

	for _ in range(self.lsh_num_pre_rounds):
	keys, values, mask = self.lsh_round(keys, values, mask, t*extra_factor)

	keys, values, mask = self.lsh_round(keys, values, mask, t//self.sparsity_factor)
	keys /= mask + 1e-8
	values /= mask + 1e-8

	mask = (1. - mask.clamp(0, 1))
	mask *= torch.finfo(mask.dtype).min

	return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.transpose(-1, -2).reshape(n, h, 1, -1)

	def lsh_round(self, keys, values, mask, output_size):

	with torch.no_grad():

	n_hashes = output_size // 2
	n, h, b, t, d = keys.size()
	binary_mask = mask.clamp(0, 1)

	indexes = (torch.nn.functional.normalize(keys, dim=-1) * binary_mask) @ torch.randn(1, h, 1, d, n_hashes, device=keys.device)
	indexes = torch.cat([indexes, -indexes], dim=-1).argmax(dim=-1, keepdim=True)

	n, h, b, t, d = keys.size()

	x_ = torch.zeros(n, h, b, output_size, d, device=keys.device)
	mask_ = torch.zeros(n, h, b, output_size, 1, device=keys.device)
	keys = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=keys)
	values = torch.scatter_add(x_, dim=-2, index=indexes.expand(-1, -1, -1, -1, d), src=values)
	mask = torch.scatter_add(mask_, dim=-2, index=indexes, src=mask)

	return keys[..., :output_size, :], values[..., :output_size, :], mask[..., :output_size, :]

	def get_sparse_tokens_with_bos_pooling(self, queries, keys, values, mask):

	if self.sparsity_factor == 1:
	return keys, values, mask.expand(-1, keys.size()[1], -1, -1)

	queries = queries.unsqueeze(-3)
	mask = self.chunk(mask.transpose(-1, -2), self.sparsity_factor).transpose(-1, -2)
	keys = self.chunk(keys, self.sparsity_factor)
	values = self.chunk(values, self.sparsity_factor)

	n, h, b, t, d = keys.size()
	scores = (queries[..., :1, :] @ keys.transpose(-1, -2)) / math.sqrt(d)
	if mask is not None:
	scores = scores + mask

	scores = torch.softmax(scores, dim=-1)
	keys = scores @ keys
	values = scores @ values
	mask = mask.mean(dim=-1)
	mask[mask != torch.finfo(mask.dtype).min] = 0

	return keys.reshape(n, h, -1, d), values.reshape(n, h, -1, d), mask.expand(-1, h, -1, -1).transpose(-1, -2)

	def forward(
	self,
	hidden_states,
	attention_mask=None,
	head_mask=None,
	encoder_hidden_states=None,
	encoder_attention_mask=None,
	past_key_value=None,
	output_attentions=False,
	):

	query_layer = self.query(hidden_states)

	# If this is instantiated as a cross-attention module, the keys
	# and values come from an encoder; the attention mask needs to be
	# such that the encoder's padding tokens are not attended to.
	is_cross_attention = encoder_hidden_states is not None

	if is_cross_attention and past_key_value is not None:
	# reuse k,v, cross_attentions
	key_layer = past_key_value[0]
	value_layer = past_key_value[1]
	attention_mask = encoder_attention_mask
	elif is_cross_attention:
	key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
	value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
	attention_mask = encoder_attention_mask
	elif past_key_value is not None:
	key_layer = self.transpose_for_scores(self.key(hidden_states))
	value_layer = self.transpose_for_scores(self.value(hidden_states))
	key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
	value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
	else:
	key_layer = self.transpose_for_scores(self.key(hidden_states))
	value_layer = self.transpose_for_scores(self.value(hidden_states))

	query_layer = self.transpose_for_scores(query_layer)

	if self.is_decoder:
	# if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
	# Further calls to cross_attention layer can then reuse all cross-attention
	# key/value_states (first "if" case)
	# if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
	# all previous decoder key/value_states. Further calls to uni-directional self-attention
	# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
	# if encoder bi-directional self-attention `past_key_value` is always `None`
	past_key_value = (key_layer, value_layer)

	if is_cross_attention:
	outputs = self.cross_attention_forward(
	query_layer=query_layer,
	key_layer=key_layer,
	value_layer=value_layer,
	attention_mask=attention_mask,
	output_attentions=output_attentions
	)
	else:
	outputs = self.causal_forward(
	query_layer,
	key_layer,
	value_layer,
	attention_mask=attention_mask,
	output_attentions=output_attentions,
	)

	outputs = outputs + ((key_layer, value_layer),)

	else:
	outputs = self.not_causal_forward(
	query_layer,
	key_layer,
	value_layer,
	attention_mask=attention_mask,
	output_attentions=output_attentions
	)

	return outputs

	def causal_forward(
	self,
	query_layer,
	key_layer,
	value_layer,
	attention_mask=None,
	output_attentions=False,
	):

	n, h, t, d = key_layer.size()

	# Cat global mask
	attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)

	# Split input into global tokens and other tokens
	split = (self.num_global_tokens, t - self.num_global_tokens)
	global_query, query_layer = query_layer.split(split, dim=-2)

	# Use normal causal attention if local attention covers every tokens
	if t <= 2 * self.block_size + self.num_global_tokens:
	context_layer = self.causal_attention(
	query_layer=query_layer,
	key_layer=key_layer,
	value_layer=value_layer,
	attention_mask=attention_mask,
	causal_shape=(t - self.num_global_tokens, t - self.num_global_tokens)
	)

	context_layer = torch.cat([global_query, context_layer], dim=-2)
	return (self.reshape_output(context_layer), )

	# Split K Q M on global and non global
	global_key, key_layer = key_layer.split(split, dim=-2)
	global_value, value_layer = value_layer.split(split, dim=-2)
	global_mask, attention_mask = attention_mask.split(split, dim=-1)

	n, h, t, d = key_layer.size()

	# Get sparse idx
	sparse_key, sparse_value, sparse_mask = (None, None, None)
	if self.sparse_block_size and self.sparsity_factor > 0:
	sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(query_layer, key_layer, value_layer, attention_mask)

	# Expand masks on heads
	attention_mask = attention_mask.expand(-1, h, -1, -1)
	global_mask = global_mask.expand(-1, h, -1, -1)

	# Compute dot product attention
	context_layer = self.attention(
	query_layer,
	key_layer,
	value_layer,
	attention_mask,
	sparse_key=sparse_key,
	sparse_value=sparse_value,
	sparse_mask=sparse_mask,
	global_key=global_key,
	global_value=global_value,
	global_mask=global_mask
	)

	# Merge pseudo global (causal) and local-sparse tokens
	context_layer = torch.cat([global_query, context_layer], dim=-2)
	context_layer = self.reshape_output(context_layer)

	return (context_layer,)

	def not_causal_forward(
	self,
	query_layer,
	key_layer,
	value_layer,
	attention_mask=None,
	output_attentions=False,
	):

	n, h, t, d = query_layer.size()

	# Cat global mask
	attention_mask = torch.nn.functional.pad(attention_mask, (self.num_global_tokens, 0), value=0)

	# Use normal attention if local attention covers every tokens
	if t <= 2 * self.block_size + self.num_global_tokens:
	context_layer = self.full_attention(
	query_layer=query_layer,
	key_layer=key_layer,
	value_layer=value_layer,
	attention_mask=attention_mask
	)
	return (self.reshape_output(context_layer), )

	# Split input into global tokens and other tokens
	split = (self.num_global_tokens, t - self.num_global_tokens)
	global_query, query_layer = query_layer.split(split, dim=-2)

	# Get global_attention
	bos = self.full_attention(
	query_layer=global_query,
	key_layer=key_layer,
	value_layer=value_layer,
	attention_mask=attention_mask
	)

	# Split K Q M on global and non global
	global_key, key_layer = key_layer.split(split, dim=-2)
	global_value, value_layer = value_layer.split(split, dim=-2)
	global_mask, attention_mask = attention_mask.split(split, dim=-1)

	n, h, t, d = key_layer.size()

	# Get sparse idx
	sparse_key, sparse_value, sparse_mask = (None, None, None)

	if self.sparse_block_size and self.sparsity_factor > 0:
	sparse_key, sparse_value, sparse_mask = self.get_sparse_elements(query_layer, key_layer, value_layer, attention_mask)

	# Expand masks on heads
	attention_mask = attention_mask.expand(-1, h, -1, -1)
	global_mask = global_mask.expand(-1, h, -1, -1)

	# Compute dot product attention
	context_layer = self.attention(
	query_layer,
	key_layer,
	value_layer,
	attention_mask,
	sparse_key=sparse_key,
	sparse_value=sparse_value,
	sparse_mask=sparse_mask,
	global_key=global_key,
	global_value=global_value,
	global_mask=global_mask
	)

	# Merge global and local-sparse tokens
	context_layer = torch.cat([bos, context_layer], dim=-2)
	context_layer = self.reshape_output(context_layer)

	return (context_layer,)

	def cross_attention_forward(
	self,
	query_layer,
	key_layer,
	value_layer,
	attention_mask=None,
	output_attentions=False,
	):

	context_layer = self.full_attention(
	query_layer=query_layer,
	key_layer=key_layer,
	value_layer=value_layer,
	attention_mask=attention_mask
	)
	return (self.reshape_output(context_layer), )

	def chunk(self, x, chunk_size):

	n, h, t, d = x.size()
	return x.reshape(n, h, -1, chunk_size, d)


	class LSGAttention(BertAttention):

	def __init__(self, config):

	nn.Module.__init__(self)

	self.self = LSGSelfAttention(config)
	self.output = BertSelfOutput(config)
	self.pruned_heads = set()


	class LSGBertLayer(BertLayer):

	def __init__(self, config):

	super().__init__(config)

	self.attention = LSGAttention(config)
	if self.add_cross_attention:
	if not self.is_decoder:
	assert self.is_decoder, f"{self} should be used as a decoder model if cross attention is added"
	self.crossattention = LSGAttention(config)


	class LSGBertEncoder(BertEncoder):

	def __init__(self, config):

	super().__init__(config)

	self.layer = nn.ModuleList([LSGBertLayer(config) for _ in range(config.num_hidden_layers)])

	assert hasattr(config, "num_global_tokens")
	self.num_global_tokens = config.num_global_tokens
	self.pad_idx = config.pad_token_id

	assert hasattr(config, "block_size") and hasattr(config, "adaptive")
	self.block_size = config.block_size
	self.adaptive = config.adaptive
	self.mask_first_token = config.mask_first_token
	self.pool_with_global = config.pool_with_global

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	encoder_hidden_states: Optional[torch.FloatTensor] = None,
	encoder_attention_mask: Optional[torch.FloatTensor] = None,
	past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
	use_cache: Optional[bool] = None,
	output_attentions: Optional[bool] = False,
	output_hidden_states: Optional[bool] = False,
	return_dict: Optional[bool] = True,
	) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:

	mask_value = torch.finfo(attention_mask.dtype).min
	n, _, __, t = attention_mask.size()

	if not (self.config.is_decoder and encoder_hidden_states is not None):

	b = self.block_size * 2
	pad = t % self.block_size

	# Check if t is multiple of block_size and pad
	if self.adaptive and t > b and pad > 0:
	pad_length = self.block_size - pad
	hidden_states = torch.nn.functional.pad(hidden_states.transpose(-1, -2), (0, pad_length), value=0.).transpose(-1, -2)
	attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=mask_value)

	if self.mask_first_token:
	attention_mask[..., 0] = mask_value

	encoder_outputs = super().forward(
	hidden_states=hidden_states,
	attention_mask=attention_mask,
	head_mask=head_mask,
	encoder_hidden_states=encoder_hidden_states,
	encoder_attention_mask=encoder_attention_mask,
	past_key_values=past_key_values,
	use_cache=use_cache,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict
	)

	sequence_output = encoder_outputs[0]
	if self.pool_with_global:
	sequence_output[:, self.num_global_tokens] = sequence_output[:, 0]

	# Adapt sequence to initial shape
	sequence_output = sequence_output[..., self.num_global_tokens: t + self.num_global_tokens, :]

	if not return_dict:
	return (sequence_output, ) + encoder_outputs[1:]

	encoder_outputs.last_hidden_state = sequence_output
	return encoder_outputs


	class LSGBertPreTrainedModel(BertPreTrainedModel):
	"""
	An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
	models.
	"""

	config_class = LSGBertConfig

	def _set_gradient_checkpointing(self, module, value=False):
	if isinstance(module, (BertEncoder, LSGBertEncoder)):
	module.gradient_checkpointing = value


	class LSGBertModel(LSGBertPreTrainedModel, BertModel):
	"""
	This class overrides :class:`~transformers.BertModel`. Please check the superclass for the appropriate
	documentation alongside usage examples.
	"""

	def __init__(self, config, add_pooling_layer=True):

	LSGBertPreTrainedModel.__init__(self, config)

	self.config = config

	self.embeddings = LSGBertEmbeddings(config)
	self.encoder = LSGBertEncoder(config)
	self.pooler = BertPooler(config) if add_pooling_layer else None

	if config.add_cross_attention:
	logger.warning(
	"Cross attention is computed using full attention since it is not LSG compatible."
	)

	self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
	if self._use_flash_attention_2:
	logger.warning(
	"[WARNING flash-attention]: LSG doesnt support flash-attention currently"
	)

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

	def get_extended_attention_mask(self, attention_mask, input_shape, device=None):

	# Do not rely on original triangular mask from BERT/RoBERTa for causalLM
	if attention_mask.dim() == 3:
	extended_attention_mask = attention_mask[:, None, :, :]
	elif attention_mask.dim() == 2:
	extended_attention_mask = attention_mask[:, None, None, :]
	else:
	raise ValueError(
	f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
	)

	extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
	extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(extended_attention_mask.dtype).min

	return extended_attention_mask


	class LSGBertForPreTraining(LSGBertPreTrainedModel, BertForPreTraining):

	_tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	self.bert = LSGBertModel(config)
	self.cls = BertPreTrainingHeads(config)

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


	class LSGBertLMHeadModel(LSGBertPreTrainedModel, BertLMHeadModel):

	_tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	if not config.is_decoder:
	logger.warning("If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`")

	self.bert = LSGBertModel(config, add_pooling_layer=False)
	self.cls = BertOnlyMLMHead(config)

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


	class LSGBertForMaskedLM(LSGBertPreTrainedModel, BertForMaskedLM):
	"""
	This class overrides :class:`~transformers.BertForMaskedLM`. Please check the superclass for the appropriate
	documentation alongside usage examples.
	"""

	_tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	if config.is_decoder:
	logger.warning(
	"If you want to use `LSGBertForMaskedLM` make sure `config.is_decoder=False` for "
	"bi-directional self-attention."
	)

	self.bert = LSGBertModel(config, add_pooling_layer=False)
	self.cls = BertOnlyMLMHead(config)

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


	class LSGBertForNextSentencePrediction(LSGBertPreTrainedModel, BertForNextSentencePrediction):

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	self.bert = LSGBertModel(config)
	self.cls = BertOnlyNSPHead(config)

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


	class LSGBertForSequenceClassification(LSGBertPreTrainedModel, BertForSequenceClassification):
	"""
	This class overrides :class:`~transformers.BertForSequenceClassification`. Please check the superclass for the
	appropriate documentation alongside usage examples.
	"""

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	self.num_labels = config.num_labels
	self.config = config

	self.bert = LSGBertModel(config)
	classifier_dropout = (
	config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
	)
	self.dropout = nn.Dropout(classifier_dropout)
	self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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


	class LSGBertForMultipleChoice(LSGBertPreTrainedModel, BertForMultipleChoice):
	"""
	This class overrides :class:`~transformers.BertForMultipleChoice`. Please check the superclass for the
	appropriate documentation alongside usage examples.
	"""

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	self.bert = LSGBertModel(config)
	classifier_dropout = (
	config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
	)
	self.dropout = nn.Dropout(classifier_dropout)
	self.classifier = nn.Linear(config.hidden_size, 1)

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


	class LSGBertForTokenClassification(LSGBertPreTrainedModel, BertForTokenClassification):
	"""
	This class overrides :class:`~transformers.BertForTokenClassification`. Please check the superclass for the
	appropriate documentation alongside usage examples.
	"""

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	self.num_labels = config.num_labels

	self.bert = LSGBertModel(config, add_pooling_layer=False)
	classifier_dropout = (
	config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
	)
	self.dropout = nn.Dropout(classifier_dropout)
	self.classifier = nn.Linear(config.hidden_size, config.num_labels)

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


	class LSGBertForQuestionAnswering(LSGBertPreTrainedModel, BertForQuestionAnswering):
	"""
	This class overrides :class:`~transformers.BertForQuestionAnswering`. Please check the superclass for the
	appropriate documentation alongside usage examples.
	"""

	def __init__(self, config):

	LSGBertPreTrainedModel.__init__(self, config)

	self.num_labels = config.num_labels

	self.bert = LSGBertModel(config, add_pooling_layer=False)
	self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

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


	def str_to_class(classname):
	return getattr(sys.modules[__name__], classname)

	# Register model in Auto API
	try:
	LSGBertConfig.register_for_auto_class()
	for key, value in AUTO_MAP.items():
	str_to_class(value.split(".")[-1]).register_for_auto_class(key)
	except:
	warn("AutoRegister isn't available, you'll have to manually copy modeling.py after .save_pretrained(...).")
	warn("Update to transformers >= 4.36.1 to fix.")