| import math
|
| from typing import Optional
|
|
|
| import torch
|
| from torch import Tensor
|
| from torch.nn import (
|
| BCEWithLogitsLoss,
|
| Conv1d,
|
| CrossEntropyLoss,
|
| Dropout,
|
| Embedding,
|
| Linear,
|
| Module,
|
| ModuleList,
|
| MSELoss,
|
| Parameter,
|
| SiLU,
|
| functional as F,
|
| )
|
| from transformers import PreTrainedModel
|
| from transformers.modeling_outputs import (
|
| BaseModelOutput,
|
| MaskedLMOutput,
|
| SequenceClassifierOutput,
|
| TokenClassifierOutput,
|
| )
|
| from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
|
|
|
| try:
|
| from .iii_configuration import IIIBERTConfig
|
| except ImportError:
|
| from iii_configuration import IIIBERTConfig
|
|
|
| _CONFIG_FOR_DOC = "IIIBERTConfig"
|
|
|
|
|
| class IIIRMSNorm(Module):
|
| def __init__(self, dim, eps=1e-6):
|
| super().__init__()
|
| self.weight = Parameter(torch.ones(dim))
|
| self.epsilon = eps
|
|
|
| def forward(self, tensor):
|
| input_dtype = tensor.dtype
|
| tensor = tensor.to(torch.float32)
|
| variance = tensor.pow(2).mean(-1, keepdim=True)
|
| tensor = tensor * torch.rsqrt(variance + self.epsilon)
|
| return self.weight * tensor.to(input_dtype)
|
|
|
|
|
| ALL_LAYERNORM_LAYERS.append(IIIRMSNorm)
|
|
|
|
|
| class IIIMLP(Module):
|
| def __init__(self, config, dim=None, intermediate_dim=None):
|
| super().__init__()
|
| self.config = config
|
| self.dim = config.hidden_dim if dim is None else dim
|
| self.intermediate_dim = (
|
| config.intermediate_dim if intermediate_dim is None else intermediate_dim
|
| )
|
|
|
| self.gate = Linear(self.dim, self.intermediate_dim, bias=False)
|
| self.up = Linear(self.dim, self.intermediate_dim, bias=False)
|
| self.down = Linear(self.intermediate_dim, self.dim, bias=False)
|
| self.activation = SiLU()
|
|
|
| def forward(self, x):
|
| value = self.up(x)
|
| gate = self.activation(self.gate(x))
|
| return self.down(gate * value)
|
|
|
|
|
| class IIIAttention(Module):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__()
|
| self.config = config
|
|
|
| self.hidden_dim = config.hidden_dim
|
| self.lora_dim = config.lora_dim
|
| self.num_heads = config.n_heads
|
|
|
| if self.lora_dim % self.num_heads != 0:
|
| raise ValueError(
|
| f"lora_dim ({self.lora_dim}) must be divisible by "
|
| f"n_heads ({self.num_heads})."
|
| )
|
|
|
| self.head_dim = self.lora_dim // self.num_heads
|
|
|
| self.qkv = Linear(self.hidden_dim, self.lora_dim * 3, bias=config.attention_bias)
|
| self.o = Linear(self.lora_dim, self.hidden_dim, bias=config.attention_bias)
|
|
|
| self.q_norm = IIIRMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
| self.k_norm = IIIRMSNorm(self.head_dim, eps=config.rms_norm_eps)
|
|
|
| def _prepare_attention_mask(self, attention_mask, device):
|
| if attention_mask is None:
|
| return None
|
|
|
| mask = attention_mask.to(device)
|
|
|
| if mask.dim() == 2:
|
| mask = mask[:, None, None, :]
|
| elif mask.dim() == 3:
|
| mask = mask[:, None, :, :]
|
| elif mask.dim() != 4:
|
| raise ValueError("attention_mask must be 2D, 3D, or 4D.")
|
|
|
| if not torch.is_floating_point(mask) and mask.dtype != torch.bool:
|
| mask = mask != 0
|
|
|
| return mask
|
|
|
| def forward(self, hidden_states, attention_mask=None, output_attentions=False):
|
| batch_size, seq_len, _ = hidden_states.shape
|
|
|
| q, k, v = self.qkv(hidden_states).chunk(3, dim=-1)
|
|
|
| q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
| k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
| v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
|
|
|
| q = self.q_norm(q)
|
| k = self.k_norm(k)
|
|
|
| mask = self._prepare_attention_mask(attention_mask, hidden_states.device)
|
|
|
| if mask is not None and mask.dtype != torch.bool:
|
| mask = mask.to(q.dtype)
|
|
|
| attention_weights = None
|
|
|
| if output_attentions:
|
| scale = self.head_dim ** -0.5
|
| scores = torch.matmul(q, k.transpose(-2, -1)) * scale
|
|
|
| if mask is not None:
|
| if mask.dtype == torch.bool:
|
| scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
|
| else:
|
| scores = scores + mask.to(scores.dtype)
|
|
|
| attention_weights = torch.softmax(scores, dim=-1)
|
| attn_output = torch.matmul(attention_weights, v)
|
| else:
|
| attn_output = F.scaled_dot_product_attention(
|
| q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
|
| )
|
|
|
| attn_output = (
|
| attn_output.transpose(1, 2)
|
| .contiguous()
|
| .view(batch_size, seq_len, self.lora_dim)
|
| )
|
|
|
| return self.o(attn_output), attention_weights
|
|
|
|
|
| class GatedCNN(Module):
|
| def __init__(self, config: IIIBERTConfig, layer_id):
|
| super().__init__()
|
| self.config = config
|
| self.hidden_dim = config.hidden_dim
|
| self.window_size = config.layers_kernel_size[layer_id]
|
| self.dilation = config.layers_diliation_size[layer_id]
|
|
|
| if self.window_size <= 0:
|
| raise ValueError("GatedCNN requires window_size > 0.")
|
| if self.dilation <= 0:
|
| raise ValueError("GatedCNN requires dilation > 0.")
|
|
|
| self.gate_proj = Linear(self.hidden_dim, self.hidden_dim * 2, bias=False)
|
|
|
| self.conv = Conv1d(
|
| in_channels=self.hidden_dim,
|
| out_channels=self.hidden_dim,
|
| kernel_size=self.window_size,
|
| dilation=self.dilation,
|
| padding="same",
|
| groups=self.hidden_dim,
|
| )
|
|
|
| self.out_proj = Linear(self.hidden_dim, self.hidden_dim)
|
|
|
| def _gate(self, hidden_states):
|
| value, score = self.gate_proj(hidden_states).chunk(2, dim=-1)
|
| return value * torch.sigmoid(score)
|
|
|
| @staticmethod
|
| def _keep_mask(attention_mask):
|
| if attention_mask is None or attention_mask.dim() != 2:
|
| return None
|
| if torch.is_floating_point(attention_mask):
|
| keep = attention_mask >= 0
|
| else:
|
| keep = attention_mask != 0
|
| return keep.unsqueeze(-1)
|
|
|
| def forward(self, hidden_states, attention_mask=None):
|
| gated = self._gate(hidden_states)
|
|
|
| keep = self._keep_mask(attention_mask)
|
| if keep is not None:
|
| gated = gated * keep.to(gated.dtype)
|
|
|
| out = gated.transpose(1, 2)
|
| out = self.conv(out)
|
| out = out.transpose(1, 2)
|
|
|
| if keep is not None:
|
| out = out * keep.to(out.dtype)
|
|
|
| return self.out_proj(out)
|
|
|
|
|
| class IIILayer(Module):
|
| def __init__(self, config: IIIBERTConfig, layer_id):
|
| super().__init__()
|
| self.config = config
|
| self.layer_id = layer_id
|
|
|
| self.window_size = config.layers_kernel_size[layer_id]
|
| self.dilation = config.layers_diliation_size[layer_id]
|
|
|
| self.input_layernorm = IIIRMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
|
|
|
| if self.window_size == 0:
|
| self.is_attention = True
|
| self.mixer = IIIAttention(config)
|
| else:
|
| self.is_attention = False
|
| self.mixer = GatedCNN(config, layer_id)
|
|
|
| self.post_mixer_layernorm = IIIRMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
|
| self.mlp = IIIMLP(config)
|
|
|
| def forward(self, hidden_states, attention_mask=None, output_attentions=False):
|
| residual = hidden_states
|
| hidden_states = self.input_layernorm(hidden_states)
|
|
|
| if self.is_attention:
|
| mixer_output, attention_weights = self.mixer(
|
| hidden_states,
|
| attention_mask=attention_mask,
|
| output_attentions=output_attentions,
|
| )
|
| else:
|
| mixer_output = self.mixer(hidden_states, attention_mask=attention_mask)
|
| attention_weights = None
|
|
|
| hidden_states = residual + mixer_output
|
|
|
| residual = hidden_states
|
| hidden_states = self.post_mixer_layernorm(hidden_states)
|
| hidden_states = self.mlp(hidden_states)
|
| hidden_states = residual + hidden_states
|
|
|
| return hidden_states, attention_weights
|
|
|
|
|
| class IIIBERTLayerGroup(Module):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__()
|
| self.layers = ModuleList(
|
| [IIILayer(config, i) for i in range(config.n_sublayers)]
|
| )
|
| self.norm = IIIRMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
|
|
|
| def forward(self, hidden_states, attention_mask=None, output_attentions=False):
|
| all_attentions = []
|
|
|
| for layer in self.layers:
|
| hidden_states, attention_weights = layer(
|
| hidden_states,
|
| attention_mask=attention_mask,
|
| output_attentions=output_attentions,
|
| )
|
| if attention_weights is not None:
|
| all_attentions.append(attention_weights)
|
|
|
| return hidden_states, tuple(all_attentions)
|
|
|
|
|
| class SinusiodalPositionalEncoding(Module):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__()
|
|
|
| self.dropout = Dropout(config.dropout)
|
|
|
| pe = torch.zeros(config.max_sequence_length, config.embedding_dim)
|
| position = torch.arange(config.max_sequence_length, dtype=torch.float32).unsqueeze(1)
|
|
|
| div_term = torch.exp(
|
| torch.arange(0, config.embedding_dim, 2, dtype=torch.float32)
|
| * (-math.log(config.pe_theta) / config.embedding_dim)
|
| )
|
|
|
| pe[:, 0::2] = torch.sin(position * div_term)
|
| pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].shape[1]])
|
|
|
| self.register_buffer("pe", pe.unsqueeze(0))
|
|
|
| def forward(self, x: Tensor):
|
| seq_len = x.size(1)
|
|
|
| if seq_len > self.pe.size(1):
|
| raise ValueError("Sequence length exceeds max_sequence_length")
|
|
|
| x = x + self.pe[:, :seq_len, :].to(dtype=x.dtype)
|
| return self.dropout(x)
|
|
|
|
|
| class IIIBERTPreTrainedModel(PreTrainedModel):
|
| config_class = IIIBERTConfig
|
| base_model_prefix = "iii_bert"
|
| supports_gradient_checkpointing = False
|
| _no_split_modules = ["IIILayer"]
|
|
|
| def _init_weights(self, module):
|
| std = self.config.initializer_range
|
|
|
| if isinstance(module, Linear):
|
| module.weight.data.normal_(mean=0.0, std=std)
|
| if module.bias is not None:
|
| module.bias.data.zero_()
|
| elif isinstance(module, Conv1d):
|
| module.weight.data.normal_(mean=0.0, std=std)
|
| if module.bias is not None:
|
| module.bias.data.zero_()
|
| elif isinstance(module, Embedding):
|
| module.weight.data.normal_(mean=0.0, std=std)
|
| if module.padding_idx is not None:
|
| module.weight.data[module.padding_idx].zero_()
|
| elif isinstance(module, IIIRMSNorm):
|
| module.weight.data.fill_(1.0)
|
|
|
|
|
| class IIIBERTModel(IIIBERTPreTrainedModel):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__(config)
|
|
|
| self.n_layers = config.n_layers
|
|
|
| self.embed_tokens = Embedding(config.vocab_size, config.embedding_dim)
|
| self.pe = SinusiodalPositionalEncoding(config)
|
| self.norm = IIIRMSNorm(config.embedding_dim, eps=config.rms_norm_eps)
|
| self.linear_proj = Linear(config.embedding_dim, config.hidden_dim, bias=False)
|
|
|
| self.model = IIIBERTLayerGroup(config)
|
|
|
| self.post_init()
|
|
|
| def get_input_embeddings(self):
|
| return self.embed_tokens
|
|
|
| def set_input_embeddings(self, value):
|
| self.embed_tokens = value
|
|
|
| def forward(
|
| self,
|
| input_ids: Optional[Tensor] = None,
|
| attention_mask: Optional[Tensor] = None,
|
| token_type_ids: Optional[Tensor] = None,
|
| inputs_embeds: Optional[Tensor] = None,
|
| output_attentions: Optional[bool] = None,
|
| output_hidden_states: Optional[bool] = None,
|
| return_dict: Optional[bool] = None,
|
| ):
|
| 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 None) == (inputs_embeds is None):
|
| raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
|
|
|
| embed = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
|
| embed = self.pe(embed)
|
| embed = self.norm(embed)
|
| embed = self.linear_proj(embed)
|
|
|
| all_hidden_states = (embed,) if output_hidden_states else None
|
| all_attentions = () if output_attentions else None
|
|
|
| for i in range(self.n_layers):
|
| embed, attentions = self.model(
|
| embed,
|
| attention_mask=attention_mask,
|
| output_attentions=output_attentions,
|
| )
|
|
|
| if i == self.n_layers - 1:
|
| embed = self.model.norm(embed)
|
|
|
| if output_hidden_states:
|
| all_hidden_states += (embed,)
|
| if output_attentions:
|
| all_attentions += attentions
|
|
|
| if not return_dict:
|
| return tuple(
|
| v for v in (embed, all_hidden_states, all_attentions) if v is not None
|
| )
|
|
|
| return BaseModelOutput(
|
| last_hidden_state=embed,
|
| hidden_states=all_hidden_states,
|
| attentions=all_attentions,
|
| )
|
|
|
|
|
| class IIIBERTForMaskedLM(IIIBERTPreTrainedModel):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__(config)
|
|
|
| self.iii_bert = IIIBERTModel(config)
|
| self.linear_out = Linear(config.hidden_dim, config.embedding_dim, bias=False)
|
|
|
| self.post_init()
|
|
|
| def forward(
|
| self,
|
| input_ids: Optional[Tensor] = None,
|
| attention_mask: Optional[Tensor] = None,
|
| token_type_ids: Optional[Tensor] = None,
|
| inputs_embeds: Optional[Tensor] = None,
|
| labels: Optional[Tensor] = None,
|
| output_attentions: Optional[bool] = None,
|
| output_hidden_states: Optional[bool] = None,
|
| return_dict: Optional[bool] = None,
|
| ):
|
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
| outputs = self.iii_bert(
|
| input_ids=input_ids,
|
| attention_mask=attention_mask,
|
| inputs_embeds=inputs_embeds,
|
| output_attentions=output_attentions,
|
| output_hidden_states=output_hidden_states,
|
| return_dict=True,
|
| )
|
|
|
| logits = F.linear(
|
| self.linear_out(outputs.last_hidden_state),
|
| self.iii_bert.embed_tokens.weight,
|
| )
|
|
|
| loss = None
|
| if labels is not None:
|
| loss = CrossEntropyLoss()(
|
| logits.view(-1, self.config.vocab_size), labels.view(-1)
|
| )
|
|
|
| if not return_dict:
|
| output = (logits,) + tuple(
|
| v for v in (outputs.hidden_states, outputs.attentions) if v is not None
|
| )
|
| return ((loss,) + output) if loss is not None else output
|
|
|
| return MaskedLMOutput(
|
| loss=loss,
|
| logits=logits,
|
| hidden_states=outputs.hidden_states,
|
| attentions=outputs.attentions,
|
| )
|
|
|
|
|
| class IIIBERTForSequenceClassification(IIIBERTPreTrainedModel):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__(config)
|
|
|
| self.num_labels = config.num_labels
|
|
|
| self.iii_bert = IIIBERTModel(config)
|
| self.dropout = Dropout(config.dropout)
|
| self.classifier = Linear(config.hidden_dim, config.num_labels)
|
|
|
| self.post_init()
|
|
|
| def forward(
|
| self,
|
| input_ids: Optional[Tensor] = None,
|
| attention_mask: Optional[Tensor] = None,
|
| token_type_ids: Optional[Tensor] = None,
|
| inputs_embeds: Optional[Tensor] = None,
|
| labels: Optional[Tensor] = None,
|
| output_attentions: Optional[bool] = None,
|
| output_hidden_states: Optional[bool] = None,
|
| return_dict: Optional[bool] = None,
|
| ):
|
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
| outputs = self.iii_bert(
|
| input_ids=input_ids,
|
| attention_mask=attention_mask,
|
| inputs_embeds=inputs_embeds,
|
| output_attentions=output_attentions,
|
| output_hidden_states=output_hidden_states,
|
| return_dict=True,
|
| )
|
|
|
| hidden = outputs.last_hidden_state
|
|
|
| if attention_mask is not None and attention_mask.dim() == 2:
|
| keep = (attention_mask != 0).unsqueeze(-1).to(hidden.dtype)
|
| pooled = (hidden * keep).sum(dim=1) / keep.sum(dim=1).clamp(min=1.0)
|
| else:
|
| pooled = hidden.mean(dim=1)
|
|
|
| logits = self.classifier(self.dropout(pooled))
|
|
|
| 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 labels.dtype in (torch.long, torch.int):
|
| self.config.problem_type = "single_label_classification"
|
| else:
|
| self.config.problem_type = "multi_label_classification"
|
|
|
| if self.config.problem_type == "regression":
|
| if self.num_labels == 1:
|
| loss = MSELoss()(logits.squeeze(), labels.squeeze())
|
| else:
|
| loss = MSELoss()(logits, labels)
|
| elif self.config.problem_type == "single_label_classification":
|
| loss = CrossEntropyLoss()(
|
| logits.view(-1, self.num_labels), labels.view(-1)
|
| )
|
| else:
|
| loss = BCEWithLogitsLoss()(logits, labels)
|
|
|
| if not return_dict:
|
| output = (logits,) + tuple(
|
| v for v in (outputs.hidden_states, outputs.attentions) if v is not None
|
| )
|
| return ((loss,) + output) if loss is not None else output
|
|
|
| return SequenceClassifierOutput(
|
| loss=loss,
|
| logits=logits,
|
| hidden_states=outputs.hidden_states,
|
| attentions=outputs.attentions,
|
| )
|
|
|
|
|
| class IIIBERTForTokenClassification(IIIBERTPreTrainedModel):
|
| def __init__(self, config: IIIBERTConfig):
|
| super().__init__(config)
|
|
|
| self.num_labels = config.num_labels
|
|
|
| self.iii_bert = IIIBERTModel(config)
|
| self.dropout = Dropout(config.dropout)
|
| self.classifier = Linear(config.hidden_dim, config.num_labels)
|
|
|
| self.post_init()
|
|
|
| def forward(
|
| self,
|
| input_ids: Optional[Tensor] = None,
|
| attention_mask: Optional[Tensor] = None,
|
| token_type_ids: Optional[Tensor] = None,
|
| inputs_embeds: Optional[Tensor] = None,
|
| labels: Optional[Tensor] = None,
|
| output_attentions: Optional[bool] = None,
|
| output_hidden_states: Optional[bool] = None,
|
| return_dict: Optional[bool] = None,
|
| ):
|
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
| outputs = self.iii_bert(
|
| input_ids=input_ids,
|
| attention_mask=attention_mask,
|
| inputs_embeds=inputs_embeds,
|
| output_attentions=output_attentions,
|
| output_hidden_states=output_hidden_states,
|
| return_dict=True,
|
| )
|
|
|
| logits = self.classifier(self.dropout(outputs.last_hidden_state))
|
|
|
| loss = None
|
| if labels is not None:
|
| loss = CrossEntropyLoss()(
|
| logits.view(-1, self.num_labels), labels.view(-1)
|
| )
|
|
|
| if not return_dict:
|
| output = (logits,) + tuple(
|
| v for v in (outputs.hidden_states, outputs.attentions) if v is not None
|
| )
|
| return ((loss,) + output) if loss is not None else output
|
|
|
| return TokenClassifierOutput(
|
| loss=loss,
|
| logits=logits,
|
| hidden_states=outputs.hidden_states,
|
| attentions=outputs.attentions,
|
| )
|
|
|
|
|
| IIIBERTConfig.register_for_auto_class()
|
| IIIBERTModel.register_for_auto_class("AutoModel")
|
| IIIBERTForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")
|
| IIIBERTForSequenceClassification.register_for_auto_class("AutoModelForSequenceClassification")
|
| IIIBERTForTokenClassification.register_for_auto_class("AutoModelForTokenClassification") |
