ruRoPEBert-e5-base-512 / modeling_rope_bert.py

Added model code

30f8826 verified 3 months ago

No virus

46.2 kB

	# coding=utf-8
	# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
	# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	# See the License for the specific language governing permissions and
	# limitations under the License.
	"""PyTorch BERT model."""


	import math
	from dataclasses import dataclass
	from typing import Optional, Tuple, Union

	import torch
	import torch.utils.checkpoint
	from torch import nn
	from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
	from transformers import PretrainedConfig
	from transformers.activations import ACT2FN
	from transformers.modeling_outputs import (
	BaseModelOutputWithPooling,
	MaskedLMOutput,
	SequenceClassifierOutput,
	)
	from transformers.modeling_utils import PreTrainedModel
	from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
	from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
	from transformers.utils import (
	ModelOutput,
	add_start_docstrings,
	add_start_docstrings_to_model_forward,
	logging,
	)

	logger = logging.get_logger(__name__)


	class RoPEBertConfig(PretrainedConfig):

	model_type = "bert"

	def __init__(
	self,
	vocab_size=30522,
	hidden_size=768,
	num_hidden_layers=12,
	num_attention_heads=12,
	intermediate_size=3072,
	hidden_act="gelu",
	pooler_type="mean",
	hidden_dropout_prob=0.1,
	attention_probs_dropout_prob=0.1,
	max_position_embeddings=512,
	type_vocab_size=2,
	initializer_range=0.02,
	layer_norm_eps=1e-12,
	pad_token_id=0,
	classifier_dropout=None,
	rope_theta=10000.0,
	rope_scaling=None,
	**kwargs,
	):
	super().__init__(pad_token_id=pad_token_id, **kwargs)

	self.vocab_size = vocab_size
	self.hidden_size = hidden_size
	self.num_hidden_layers = num_hidden_layers
	self.num_attention_heads = num_attention_heads
	self.hidden_act = hidden_act
	self.intermediate_size = intermediate_size
	self.hidden_dropout_prob = hidden_dropout_prob
	self.attention_probs_dropout_prob = attention_probs_dropout_prob
	self.max_position_embeddings = max_position_embeddings
	self.type_vocab_size = type_vocab_size
	self.initializer_range = initializer_range
	self.layer_norm_eps = layer_norm_eps
	self.classifier_dropout = classifier_dropout
	self.rope_theta = rope_theta
	self.rope_scaling = rope_scaling
	self.pooler_type = pooler_type

	self._pooler_tyoe_validation()
	self._rope_scaling_validation()

	def _pooler_tyoe_validation(self):
	if self.pooler_type not in ['first_token_transform', 'mean']:
	raise ValueError(
	f"`pooler_type` must be one of `first_token_transform` or `mean`, got {self.pooler_type}"
	)

	def _rope_scaling_validation(self):
	"""
	Validate the `rope_scaling` configuration.
	"""
	if self.rope_scaling is None:
	return

	if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
	raise ValueError(
	"`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
	f"got {self.rope_scaling}"
	)
	rope_scaling_type = self.rope_scaling.get("type", None)
	rope_scaling_factor = self.rope_scaling.get("factor", None)
	if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
	raise ValueError(
	f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
	)
	if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
	raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")


	class RoPEBertEmbeddings(nn.Module):
	"""Construct the embeddings from word, token_type embeddings."""

	def __init__(self, config):
	super().__init__()
	self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
	self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

	self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	token_type_ids: Optional[torch.LongTensor] = None,
	inputs_embeds: Optional[torch.FloatTensor] = None,
	) -> torch.Tensor:
	if inputs_embeds is None:
	inputs_embeds = self.word_embeddings(input_ids)
	token_type_embeddings = self.token_type_embeddings(token_type_ids)

	embeddings = inputs_embeds + token_type_embeddings

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


	class BertRotaryEmbedding(nn.Module):
	def __init__(self, dim, max_position_embeddings=512, base=10000, device=None):
	super().__init__()

	self.dim = dim
	self.max_position_embeddings = max_position_embeddings
	self.base = base
	inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
	self.register_buffer("inv_freq", inv_freq, persistent=False)

	# Build here to make `torch.jit.trace` work.
	self._set_cos_sin_cache(
	seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
	)

	def _set_cos_sin_cache(self, seq_len, device, dtype):
	self.max_seq_len_cached = seq_len
	t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

	freqs = torch.einsum("i,j->ij", t, self.inv_freq)
	# Different from paper, but it uses a different permutation in order to obtain the same calculation
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
	self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

	def forward(self, x, seq_len=None):
	# x: [bs, num_attention_heads, seq_len, head_size]
	if seq_len > self.max_seq_len_cached:
	self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

	return (
	self.cos_cached[:seq_len].to(dtype=x.dtype),
	self.sin_cached[:seq_len].to(dtype=x.dtype),
	)


	class BertLinearScalingRotaryEmbedding(BertRotaryEmbedding):
	"""BertRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""

	def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0):
	self.scaling_factor = scaling_factor
	super().__init__(dim, max_position_embeddings, base, device)

	def _set_cos_sin_cache(self, seq_len, device, dtype):
	self.max_seq_len_cached = seq_len
	t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
	t = t / self.scaling_factor

	freqs = torch.einsum("i,j->ij", t, self.inv_freq)
	# Different from paper, but it uses a different permutation in order to obtain the same calculation
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
	self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)


	class BertDynamicNTKScalingRotaryEmbedding(BertRotaryEmbedding):
	"""BertRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""

	def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0):
	self.scaling_factor = scaling_factor
	super().__init__(dim, max_position_embeddings, base, device)

	def _set_cos_sin_cache(self, seq_len, device, dtype):
	self.max_seq_len_cached = seq_len

	if seq_len > self.max_position_embeddings:
	base = self.base * (
	(self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
	) ** (self.dim / (self.dim - 2))
	inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
	self.register_buffer("inv_freq", inv_freq, persistent=False)

	t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

	freqs = torch.einsum("i,j->ij", t, self.inv_freq)
	# Different from paper, but it uses a different permutation in order to obtain the same calculation
	emb = torch.cat((freqs, freqs), dim=-1)
	self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
	self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)


	def rotate_half(x):
	"""Rotates half the hidden dims of the input."""
	x1 = x[..., : x.shape[-1] // 2]
	x2 = x[..., x.shape[-1] // 2 :]
	return torch.cat((-x2, x1), dim=-1)


	def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
	"""Applies Rotary Position Embedding to the query and key tensors.

	Args:
	q (`torch.Tensor`): The query tensor.
	k (`torch.Tensor`): The key tensor.
	cos (`torch.Tensor`): The cosine part of the rotary embedding.
	sin (`torch.Tensor`): The sine part of the rotary embedding.
	position_ids (`torch.Tensor`):
	The position indices of the tokens corresponding to the query and key tensors. For example, this can be
	used to pass offsetted position ids when working with a KV-cache.
	unsqueeze_dim (`int`, optional, defaults to 1):
	The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
	sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
	that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
	k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
	cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
	the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
	Returns:
	`tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
	"""
	cos = cos[position_ids].unsqueeze(unsqueeze_dim)
	sin = sin[position_ids].unsqueeze(unsqueeze_dim)
	q_embed = (q * cos) + (rotate_half(q) * sin)
	k_embed = (k * cos) + (rotate_half(k) * sin)
	return q_embed, k_embed


	class RoPEBertSelfAttention(nn.Module):

	def __init__(self, config: RoPEBertConfig):
	super().__init__()
	if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
	raise ValueError(
	f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
	f"heads ({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.max_position_embeddings = config.max_position_embeddings
	self.rope_theta = config.rope_theta

	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)

	self.config = config

	self._init_rope()

	def _init_rope(self):
	if self.config.rope_scaling is None:
	self.rotary_emb = BertRotaryEmbedding(
	self.attention_head_size,
	max_position_embeddings=self.max_position_embeddings,
	base=self.rope_theta,
	)
	else:
	scaling_type = self.config.rope_scaling["type"]
	scaling_factor = self.config.rope_scaling["factor"]
	if scaling_type == "linear":
	self.rotary_emb = BertLinearScalingRotaryEmbedding(
	self.attention_head_size,
	max_position_embeddings=self.max_position_embeddings,
	scaling_factor=scaling_factor,
	base=self.rope_theta,
	)
	elif scaling_type == "dynamic":
	self.rotary_emb = BertDynamicNTKScalingRotaryEmbedding(
	self.attention_head_size,
	max_position_embeddings=self.max_position_embeddings,
	scaling_factor=scaling_factor,
	base=self.rope_theta,
	)
	else:
	raise ValueError(f"Unknown RoPE scaling type {scaling_type}")

	def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
	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 forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor]:
	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))

	kv_seq_len = key_layer.shape[-2]
	cos, sin = self.rotary_emb(value_layer, seq_len=kv_seq_len)
	query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)

	# Take the dot product between "query" and "key" to get the raw attention scores.
	attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

	attention_scores = attention_scores / math.sqrt(self.attention_head_size)
	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

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

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

	# Mask heads if we want to
	if head_mask is not None:
	attention_probs = attention_probs * head_mask

	context_layer = torch.matmul(attention_probs, value_layer)

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

	outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

	return outputs


	class RoPEBertSdpaAttention(RoPEBertSelfAttention):

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor]:

	bsz, q_len, _ = hidden_states.size()

	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))

	kv_seq_len = key_layer.shape[-2]
	cos, sin = self.rotary_emb(value_layer, seq_len=kv_seq_len)
	query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)

	# SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
	# Reference: https://github.com/pytorch/pytorch/issues/112577.
	if query_layer.device.type == "cuda" and attention_mask is not None:
	query_layer = query_layer.contiguous()
	key_layer = key_layer.contiguous()
	value_layer = value_layer.contiguous()

	context_layer = torch.nn.functional.scaled_dot_product_attention(
	query_layer,
	key_layer,
	value_layer,
	attn_mask=attention_mask,
	dropout_p=self.config.attention_probs_dropout_prob if self.training else 0.0,
	is_causal=False
	)

	context_layer = context_layer.transpose(1, 2).contiguous()
	context_layer = context_layer.reshape(bsz, q_len, self.all_head_size)

	outputs = (context_layer,)

	return outputs


	ROPEBERT_ATTENTION_CLASSES = {
	"eager": RoPEBertSelfAttention,
	"sdpa": RoPEBertSdpaAttention,
	}


	class RoPEBertSelfOutput(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.hidden_size)
	self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.dropout(hidden_states)
	hidden_states = self.LayerNorm(hidden_states + input_tensor)
	return hidden_states


	class RoPEBertAttention(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.self = ROPEBERT_ATTENTION_CLASSES[config._attn_implementation](config=config)
	self.output = RoPEBertSelfOutput(config)
	self.pruned_heads = set()

	def prune_heads(self, heads):
	if len(heads) == 0:
	return
	heads, index = find_pruneable_heads_and_indices(
	heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
	)

	# Prune linear layers
	self.self.query = prune_linear_layer(self.self.query, index)
	self.self.key = prune_linear_layer(self.self.key, index)
	self.self.value = prune_linear_layer(self.self.value, index)
	self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

	# Update hyper params and store pruned heads
	self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
	self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
	self.pruned_heads = self.pruned_heads.union(heads)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor]:
	self_outputs = self.self(
	hidden_states,
	attention_mask,
	head_mask,
	position_ids,
	output_attentions
	)
	attention_output = self.output(self_outputs[0], hidden_states)
	outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
	return outputs


	class RoPEBertIntermediate(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
	if isinstance(config.hidden_act, str):
	self.intermediate_act_fn = ACT2FN[config.hidden_act]
	else:
	self.intermediate_act_fn = config.hidden_act

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.intermediate_act_fn(hidden_states)
	return hidden_states


	class RoPEBertOutput(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
	self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.dropout = nn.Dropout(config.hidden_dropout_prob)

	def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.dropout(hidden_states)
	hidden_states = self.LayerNorm(hidden_states + input_tensor)
	return hidden_states


	class RoPEBertLayer(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.chunk_size_feed_forward = config.chunk_size_feed_forward
	self.seq_len_dim = 1
	self.attention = RoPEBertAttention(config)
	self.intermediate = RoPEBertIntermediate(config)
	self.output = RoPEBertOutput(config)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor]:
	self_attention_outputs = self.attention(
	hidden_states,
	attention_mask,
	head_mask,
	position_ids,
	output_attentions=output_attentions
	)
	attention_output = self_attention_outputs[0]
	outputs = self_attention_outputs[1:] # add self attentions if we output attention weights

	layer_output = apply_chunking_to_forward(
	self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
	)
	outputs = (layer_output,) + outputs

	return outputs

	def feed_forward_chunk(self, attention_output):
	intermediate_output = self.intermediate(attention_output)
	layer_output = self.output(intermediate_output, attention_output)
	return layer_output


	class RoPEBertEncoder(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.config = config
	self.layer = nn.ModuleList([RoPEBertLayer(config) for _ in range(config.num_hidden_layers)])
	self.gradient_checkpointing = False

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.FloatTensor] = None,
	head_mask: Optional[torch.FloatTensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	output_attentions: Optional[bool] = False,
	output_hidden_states: Optional[bool] = False,
	return_dict: Optional[bool] = True,
	) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
	all_hidden_states = () if output_hidden_states else None
	all_self_attentions = () if output_attentions else None

	for i, layer_module in enumerate(self.layer):
	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)

	layer_head_mask = head_mask[i] if head_mask is not None else None

	if self.gradient_checkpointing and self.training:
	layer_outputs = self._gradient_checkpointing_func(
	layer_module.__call__,
	hidden_states,
	attention_mask,
	layer_head_mask,
	position_ids,
	output_attentions
	)
	else:
	layer_outputs = layer_module(
	hidden_states,
	attention_mask,
	layer_head_mask,
	position_ids,
	output_attentions
	)

	hidden_states = layer_outputs[0]
	if output_attentions:
	all_self_attentions = all_self_attentions + (layer_outputs[1],)

	if output_hidden_states:
	all_hidden_states = all_hidden_states + (hidden_states,)

	if not return_dict:
	return tuple(
	v
	for v in [
	hidden_states,
	all_hidden_states,
	all_self_attentions,
	]
	if v is not None
	)
	return BaseModelOutputWithPooling(
	last_hidden_state=hidden_states,
	hidden_states=all_hidden_states,
	attentions=all_self_attentions,
	)


	class RoPEBertPooler(nn.Module):
	def __init__(self, config):
	self.pooler_type = config.pooler_type
	super().__init__()

	def forward(self, hidden_states: torch.Tensor, attention_mask: torch.LongTensor) -> torch.Tensor:
	pass


	class RoPEBertMeanTokensPooler(nn.Module):
	def __init__(self, config):
	super().__init__()

	def forward(self, hidden_states: torch.Tensor, attention_mask: torch.LongTensor) -> torch.Tensor:
	input_mask_expanded = attention_mask.unsqueeze(-1).expand(hidden_states.size()).float()
	pooled_output = torch.sum(hidden_states * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

	return pooled_output


	class RoPEBertCLSTokenTransformPooler(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.hidden_size)
	self.activation = nn.Tanh()

	def forward(self, hidden_states: torch.Tensor, attention_mask: torch.LongTensor) -> torch.Tensor:

	first_token_tensor = hidden_states[:, 0]
	pooled_output = self.dense(first_token_tensor)
	pooled_output = self.activation(pooled_output)

	return pooled_output


	ROPEBERT_POOLER_CLASSES = {
	"mean": RoPEBertMeanTokensPooler,
	"first_token_transform": RoPEBertCLSTokenTransformPooler,
	}


	class RoPEBertPredictionHeadTransform(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.dense = nn.Linear(config.hidden_size, config.hidden_size)
	if isinstance(config.hidden_act, str):
	self.transform_act_fn = ACT2FN[config.hidden_act]
	else:
	self.transform_act_fn = config.hidden_act
	self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

	def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
	hidden_states = self.dense(hidden_states)
	hidden_states = self.transform_act_fn(hidden_states)
	hidden_states = self.LayerNorm(hidden_states)
	return hidden_states


	class RoPEBertLMPredictionHead(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.transform = RoPEBertPredictionHeadTransform(config)

	# The output weights are the same as the input embeddings, but there is
	# an output-only bias for each token.
	self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

	self.bias = nn.Parameter(torch.zeros(config.vocab_size))

	# Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
	self.decoder.bias = self.bias

	def forward(self, hidden_states):
	hidden_states = self.transform(hidden_states)
	hidden_states = self.decoder(hidden_states)
	return hidden_states


	class RoPEBertOnlyMLMHead(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.predictions = RoPEBertLMPredictionHead(config)

	def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
	prediction_scores = self.predictions(sequence_output)
	return prediction_scores


	class RoPEBertOnlyNSPHead(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.seq_relationship = nn.Linear(config.hidden_size, 2)

	def forward(self, pooled_output):
	seq_relationship_score = self.seq_relationship(pooled_output)
	return seq_relationship_score


	class RoPEBertPreTrainingHeads(nn.Module):
	def __init__(self, config):
	super().__init__()
	self.predictions = RoPEBertLMPredictionHead(config)
	self.seq_relationship = nn.Linear(config.hidden_size, 2)

	def forward(self, sequence_output, pooled_output):
	prediction_scores = self.predictions(sequence_output)
	seq_relationship_score = self.seq_relationship(pooled_output)
	return prediction_scores, seq_relationship_score


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

	config_class = RoPEBertConfig
	base_model_prefix = "bert"
	supports_gradient_checkpointing = True
	_supports_sdpa = True

	def _init_weights(self, module):
	"""Initialize the weights"""
	if isinstance(module, nn.Linear):
	# Slightly different from the TF version which uses truncated_normal for initialization
	# cf https://github.com/pytorch/pytorch/pull/5617
	module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
	if module.bias is not None:
	module.bias.data.zero_()
	elif isinstance(module, nn.Embedding):
	module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
	if module.padding_idx is not None:
	module.weight.data[module.padding_idx].zero_()
	elif isinstance(module, nn.LayerNorm):
	module.bias.data.zero_()
	module.weight.data.fill_(1.0)


	@dataclass
	class RoPEBertForPreTrainingOutput(ModelOutput):

	loss: Optional[torch.FloatTensor] = None
	prediction_logits: torch.FloatTensor = None
	seq_relationship_logits: torch.FloatTensor = None
	hidden_states: Optional[Tuple[torch.FloatTensor]] = None
	attentions: Optional[Tuple[torch.FloatTensor]] = None


	class RoPEBertModel(RoPEBertPreTrainedModel):

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

	self.embeddings = RoPEBertEmbeddings(config)
	self.encoder = RoPEBertEncoder(config)

	self.pooler = ROPEBERT_POOLER_CLASSES[config.pooler_type](config=config) if add_pooling_layer else None

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

	def get_input_embeddings(self):
	return self.embeddings.word_embeddings

	def set_input_embeddings(self, value):
	self.embeddings.word_embeddings = value

	def _prune_heads(self, heads_to_prune):
	"""
	Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
	class PreTrainedModel
	"""
	for layer, heads in heads_to_prune.items():
	self.encoder.layer[layer].attention.prune_heads(heads)

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	token_type_ids: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:

	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:
	raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
	elif input_ids is not None:
	self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
	input_shape = input_ids.size()
	elif inputs_embeds is not None:
	input_shape = inputs_embeds.size()[:-1]
	else:
	raise ValueError("You have to specify either input_ids or inputs_embeds")

	if output_attentions and self.config.attn_implementation == 'sdpa':
	logger.warning("Cant use output_attentions with sdpa attention, turning off")
	output_attentions = False

	batch_size, seq_length = input_shape
	device = input_ids.device if input_ids is not None else inputs_embeds.device

	if attention_mask is None:
	attention_mask = torch.ones(((batch_size, seq_length)), device=device)

	if position_ids is None:
	position_ids = torch.arange(
	0, seq_length, dtype=torch.long, device=device
	)
	position_ids = position_ids.unsqueeze(0)

	if token_type_ids is None:
	token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

	# We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
	# ourselves in which case we just need to make it broadcastable to all heads.
	extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)

	# Prepare head mask if needed
	# 1.0 in head_mask indicate we keep the head
	# attention_probs has shape bsz x n_heads x N x N
	# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
	# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
	head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

	embedding_output = self.embeddings(
	input_ids=input_ids,
	token_type_ids=token_type_ids,
	inputs_embeds=inputs_embeds
	)
	encoder_outputs = self.encoder(
	embedding_output,
	attention_mask=extended_attention_mask,
	head_mask=head_mask,
	position_ids=position_ids,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)
	sequence_output = encoder_outputs[0]
	pooled_output = self.pooler(sequence_output, attention_mask) if self.pooler is not None else None

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

	return BaseModelOutputWithPooling(
	last_hidden_state=sequence_output,
	pooler_output=pooled_output,
	hidden_states=encoder_outputs.hidden_states,
	attentions=encoder_outputs.attentions,
	)


	class RoPEBertForPreTraining(RoPEBertPreTrainedModel):
	_tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]

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

	self.bert = RoPEBertModel(config)
	self.cls = RoPEBertPreTrainingHeads(config)

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

	def get_output_embeddings(self):
	return self.cls.predictions.decoder

	def set_output_embeddings(self, new_embeddings):
	self.cls.predictions.decoder = new_embeddings

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	token_type_ids: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.Tensor] = None,
	next_sentence_label: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple[torch.Tensor], RoPEBertForPreTrainingOutput]:

	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	outputs = self.bert(
	input_ids,
	attention_mask=attention_mask,
	token_type_ids=token_type_ids,
	position_ids=position_ids,
	head_mask=head_mask,
	inputs_embeds=inputs_embeds,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	sequence_output, pooled_output = outputs[:2]
	prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

	total_loss = None
	if labels is not None and next_sentence_label is not None:
	loss_fct = CrossEntropyLoss()
	masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
	next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
	total_loss = masked_lm_loss + next_sentence_loss

	if not return_dict:
	output = (prediction_scores, seq_relationship_score) + outputs[2:]
	return ((total_loss,) + output) if total_loss is not None else output

	return RoPEBertForPreTrainingOutput(
	loss=total_loss,
	prediction_logits=prediction_scores,
	seq_relationship_logits=seq_relationship_score,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)


	class RoPEBertForMaskedLM(RoPEBertPreTrainedModel):
	_tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]

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

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

	self.bert = RoPEBertModel(config, add_pooling_layer=False)
	self.cls = RoPEBertOnlyMLMHead(config)

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

	def get_output_embeddings(self):
	return self.cls.predictions.decoder

	def set_output_embeddings(self, new_embeddings):
	self.cls.predictions.decoder = new_embeddings

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	token_type_ids: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
	r"""
	labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, optional):
	Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
	config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
	loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
	"""

	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	outputs = self.bert(
	input_ids,
	attention_mask=attention_mask,
	token_type_ids=token_type_ids,
	position_ids=position_ids,
	head_mask=head_mask,
	inputs_embeds=inputs_embeds,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	sequence_output = outputs[0]
	prediction_scores = self.cls(sequence_output)

	masked_lm_loss = None
	if labels is not None:
	loss_fct = CrossEntropyLoss() # -100 index = padding token
	masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))

	if not return_dict:
	output = (prediction_scores,) + outputs[2:]
	return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

	return MaskedLMOutput(
	loss=masked_lm_loss,
	logits=prediction_scores,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)

	def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
	input_shape = input_ids.shape
	effective_batch_size = input_shape[0]

	# add a dummy token
	if self.config.pad_token_id is None:
	raise ValueError("The PAD token should be defined for generation")

	attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1)
	dummy_token = torch.full(
	(effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device
	)
	input_ids = torch.cat([input_ids, dummy_token], dim=1)

	return {"input_ids": input_ids, "attention_mask": attention_mask}


	class RoPEBertForSequenceClassification(RoPEBertPreTrainedModel):
	def __init__(self, config):
	super().__init__(config)
	self.num_labels = config.num_labels
	self.config = config

	self.bert = RoPEBertModel(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()

	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	token_type_ids: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	head_mask: Optional[torch.Tensor] = None,
	inputs_embeds: Optional[torch.Tensor] = None,
	labels: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
	r"""
	labels (`torch.LongTensor` of shape `(batch_size,)`, optional):
	Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
	config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
	`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
	"""
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	outputs = self.bert(
	input_ids,
	attention_mask=attention_mask,
	token_type_ids=token_type_ids,
	position_ids=position_ids,
	head_mask=head_mask,
	inputs_embeds=inputs_embeds,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	pooled_output = outputs[1]

	pooled_output = self.dropout(pooled_output)
	logits = self.classifier(pooled_output)

	loss = None
	if labels is not None:
	if self.config.problem_type is None:
	if self.num_labels == 1:
	self.config.problem_type = "regression"
	elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
	self.config.problem_type = "single_label_classification"
	else:
	self.config.problem_type = "multi_label_classification"

	if self.config.problem_type == "regression":
	loss_fct = MSELoss()
	if self.num_labels == 1:
	loss = loss_fct(logits.squeeze(), labels.squeeze())
	else:
	loss = loss_fct(logits, labels)
	elif self.config.problem_type == "single_label_classification":
	loss_fct = CrossEntropyLoss()
	loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
	elif self.config.problem_type == "multi_label_classification":
	loss_fct = BCEWithLogitsLoss()
	loss = loss_fct(logits, labels)
	if not return_dict:
	output = (logits,) + outputs[2:]
	return ((loss,) + output) if loss is not None else output

	return SequenceClassifierOutput(
	loss=loss,
	logits=logits,
	hidden_states=outputs.hidden_states,
	attentions=outputs.attentions,
	)