252 lines
10 KiB
Python
252 lines
10 KiB
Python
from typing import Any, Dict, Tuple, Union, Callable, Optional
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch
|
|
import torch.nn as nn
|
|
from typing import Tuple, Union, Callable
|
|
|
|
|
|
class EmbeddingTransformerLinear(nn.Module):
|
|
"""
|
|
A Transformer encoder-based classification model with a linear output layer.
|
|
|
|
This module uses a standard PyTorch TransformerEncoder, followed by masked mean pooling
|
|
and a linear layer applied to the pooled representation.
|
|
|
|
Args:
|
|
num_layers (int): Number of encoder layers.
|
|
vocabulary_size (int): Vocabulary size of the tokens.
|
|
hidden_size (int): Hidden size (d_model) of the Transformer.
|
|
nhead (int): Number of attention heads.
|
|
dim_feedforward (int): Feedforward dimension.
|
|
dropout (float): Dropout rate.
|
|
nonlinearity (str): Activation function ('relu' or 'gelu').
|
|
layer_norm_eps (float): Epsilon for layer normalization.
|
|
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
|
|
norm_first (bool): Whether normalization comes before sublayers.
|
|
bias (bool): Use bias in Transformer layers.
|
|
max_length (int): The maximum sequence length used for allocating the positional embedding. This value is not passed to the Transformer itself.
|
|
linear_output_size (int): Size of linear layer output.
|
|
linear_bias (bool): Use bias in linear layer.
|
|
|
|
pragma: skip_doc
|
|
"""
|
|
|
|
def __init__(self,
|
|
num_encoder_layers: int,
|
|
vocabulary_size: int,
|
|
d_model: int,
|
|
nhead: int,
|
|
dim_feedforward: int,
|
|
dropout: float,
|
|
activation: Union[str, Callable[[torch.Tensor], torch.Tensor]],
|
|
layer_norm_eps: float,
|
|
batch_first: bool,
|
|
norm_first: bool,
|
|
bias: bool,
|
|
max_length: int,
|
|
linear_output_size: int,
|
|
linear_bias: bool
|
|
):
|
|
super().__init__()
|
|
self.batch_first = batch_first
|
|
|
|
self.word_embedding = nn.Embedding(vocabulary_size, d_model)
|
|
self.positional_embedding = nn.Embedding(max_length, d_model)
|
|
|
|
# FIX: Allocate position_ids buffer once to avoid reallocation in each forward pass
|
|
# self.register_buffer("position_ids_buffer", torch.arange(max_length)) # (max_length,)
|
|
|
|
# FIX: use preallocated position_ids_buffer
|
|
# position_ids = self.position_ids_buffer[:seq_len].unsqueeze(0).expand(batch_size, -1) # (B, T)
|
|
self.position_ids = torch.arange(max_length)
|
|
|
|
encoder_layer = nn.TransformerEncoderLayer(
|
|
d_model=d_model,
|
|
nhead=nhead,
|
|
dim_feedforward=dim_feedforward,
|
|
dropout=dropout,
|
|
activation=activation,
|
|
layer_norm_eps=layer_norm_eps,
|
|
batch_first=batch_first,
|
|
norm_first=norm_first,
|
|
bias=bias
|
|
)
|
|
|
|
self.encoder = nn.TransformerEncoder(
|
|
encoder_layer=encoder_layer,
|
|
num_layers=num_encoder_layers
|
|
)
|
|
|
|
self.linear_model = nn.Linear(
|
|
in_features=d_model,
|
|
out_features=linear_output_size,
|
|
bias=linear_bias
|
|
)
|
|
torch.nn.init.xavier_uniform_(self.linear_model.weight)
|
|
|
|
def forward(self, inputs: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
|
|
"""
|
|
Forward pass through the Transformer encoder and linear layer.
|
|
|
|
Args:
|
|
inputs (Tuple[torch.Tensor, torch.Tensor]):
|
|
tokens: (B, T) token IDs,
|
|
masks: (B, T) attention mask (1 = keep, 0 = pad)
|
|
|
|
Returns:
|
|
torch.Tensor: Output of the linear layer (B, linear_output_size)
|
|
"""
|
|
tokens, masks = inputs # (B, Tokens), (B, Tokens)
|
|
if tokens.dim() != 2:
|
|
raise ValueError(f"tokens with rank {tokens.dim()} was passed.")
|
|
if masks.dim() != 2:
|
|
raise ValueError(f"masks with rank {tokens.dim()} was passed.")
|
|
if tokens.size() != masks.size():
|
|
raise ValueError(f"tokens size masks size do not match.")
|
|
|
|
batch_size, seq_len = tokens.size()
|
|
token_embeds = self.word_embedding(tokens)
|
|
|
|
position_ids_repeated = self.position_ids[:seq_len].repeat(batch_size, 1).to(token_embeds.device)
|
|
positional_embeds = self.positional_embedding(position_ids_repeated)
|
|
|
|
if token_embeds.size() != positional_embeds.size():
|
|
raise ValueError("token_embeds & positional_embeds size mismatch.")
|
|
x = token_embeds + positional_embeds # (B, Seq, Feature)
|
|
|
|
if not self.batch_first:
|
|
x = x.transpose(0, 1) # (Seq, B, Feature)
|
|
|
|
x = self.encoder(
|
|
x,
|
|
src_key_padding_mask=~masks.bool() # (B, Seq)
|
|
)
|
|
|
|
if not self.batch_first:
|
|
x = x.transpose(0, 1) # (B, Seq, Feature)
|
|
|
|
# FIX: More memory-efficient masked mean pooling
|
|
mask = masks.float() # (B, Seq)
|
|
x = (x * mask.unsqueeze(-1)).sum(dim=1) / mask.sum(dim=1, keepdim=True).clamp(min=1e-9) # (B, Feature)
|
|
|
|
outputs = self.linear_model(x) # (B, linear_output_size)
|
|
return outputs
|
|
|
|
|
|
class PtnEmbeddingTransformerLinear:
|
|
"""
|
|
A Transformer encoder model with a linear head.
|
|
|
|
Args:
|
|
num_layers (int): Number of encoder layers in Transformer.
|
|
vocabulary_size (int): The vocabulary size of the tokens.
|
|
hidden_size (int): Hidden dimension (d_model) for each token of the Transformer.
|
|
nhead (int): Number of attention heads.
|
|
dim_feedforward (int): Feedforward dimension.
|
|
dropout (float): Dropout rate for Transformer layers.
|
|
nonlinearity (str): Activation function ('relu' or 'gelu') used in Transformer layers.
|
|
layer_norm_eps (float): Epsilon value for layer normalization.
|
|
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
|
|
norm_first (bool): Whether to apply normalization before attention/feedforward.
|
|
bias (bool): Use bias in the Transformer layers.
|
|
max_length (int): Maximum input sequence length (for positional embedding).
|
|
linear_output_size (int): Output size of the final linear layer.
|
|
linear_bias (bool): Use bias in the linear layer.
|
|
|
|
category: PyTorch wrapper - Model
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the input types for the function.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary specifying required input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"num_layers": ("INT", {"default": 6, "min": 1, "max": 256}),
|
|
"vocabulary_size": ("INT", {"default": 10000, "min": 1, "max": 1e6}),
|
|
"hidden_size": ("INT", {"default": 512, "min": 1, "max": 1e6}),
|
|
"nhead": ("INT", {"default": 8, "min": 1, "max": 32}),
|
|
"dim_feedforward": ("INT", {"default": 2048, "min": 1, "max": 1e6}),
|
|
"dropout": ("FLOAT", {"default": 0.1, "min": 0, "max": 1}),
|
|
"nonlinearity": (("gelu", "relu"),),
|
|
"layer_norm_eps": ("FLOAT", {"default": 1e-5, "min": 1e-9, "max": 1e-1, "step": 1e-9}),
|
|
"batch_first": ("BOOLEAN", {"default": True}),
|
|
"norm_first": ("BOOLEAN", {"default": False}),
|
|
"bias": ("BOOLEAN", {"default": True}),
|
|
"max_length": ("INT", {"default": 512, "min": 1, "max": 1e6}),
|
|
"linear_output_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
|
|
"linear_bias": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Training"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kw):
|
|
return float("NaN")
|
|
|
|
def f(self,
|
|
num_layers: int,
|
|
vocabulary_size: int,
|
|
hidden_size: int,
|
|
nhead: int,
|
|
dim_feedforward: int,
|
|
dropout: float,
|
|
nonlinearity: str,
|
|
layer_norm_eps: float,
|
|
batch_first: bool,
|
|
norm_first: bool,
|
|
bias: bool,
|
|
max_length: int,
|
|
linear_output_size: int,
|
|
linear_bias: bool
|
|
) -> Tuple[nn.Module]:
|
|
"""
|
|
Constructs a Transformer model with a linear output layer.
|
|
|
|
Args:
|
|
num_layers (int): Number of encoder layers.
|
|
vocabulary_size (int): Vocabulary size of the tokens.
|
|
hidden_size (int): Hidden size (d_model) of the Transformer.
|
|
nhead (int): Number of attention heads.
|
|
dim_feedforward (int): Feedforward dimension.
|
|
dropout (float): Dropout rate.
|
|
nonlinearity (str): Activation function ('relu' or 'gelu').
|
|
layer_norm_eps (float): Epsilon for layer normalization.
|
|
batch_first (bool): Input shape format: True for (B, Seq, Feature), False for (Seq, B, Feature).
|
|
norm_first (bool): Whether normalization comes before sublayers.
|
|
bias (bool): Use bias in Transformer layers.
|
|
max_length (int): Maximum sequence length.
|
|
linear_output_size (int): Size of linear layer output.
|
|
linear_bias (bool): Use bias in linear layer.
|
|
|
|
Returns:
|
|
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
|
|
"""
|
|
with torch.inference_mode(False):
|
|
model = EmbeddingTransformerLinear(
|
|
num_encoder_layers=num_layers,
|
|
vocabulary_size=vocabulary_size,
|
|
d_model=hidden_size,
|
|
nhead=nhead,
|
|
dim_feedforward=dim_feedforward,
|
|
dropout=dropout,
|
|
activation=nonlinearity,
|
|
layer_norm_eps=layer_norm_eps,
|
|
batch_first=batch_first,
|
|
norm_first=norm_first,
|
|
bias=bias,
|
|
max_length=max_length,
|
|
linear_output_size=linear_output_size,
|
|
linear_bias=linear_bias
|
|
)
|
|
return (model,)
|
|
|