backend/nn/t5.py

import torch
import math

from backend.attention import attention_pytorch as attention_function


activations = {
    "gelu_pytorch_tanh": lambda a: torch.nn.functional.gelu(a, approximate="tanh"),
    "relu": torch.nn.functional.relu,
}


class T5LayerNorm(torch.nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.empty(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x):
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight.to(x) * x


class T5DenseActDense(torch.nn.Module):
    def __init__(self, model_dim, ff_dim, ff_activation):
        super().__init__()
        self.wi = torch.nn.Linear(model_dim, ff_dim, bias=False)
        self.wo = torch.nn.Linear(ff_dim, model_dim, bias=False)
        self.act = activations[ff_activation]

    def forward(self, x):
        x = self.act(self.wi(x))
        x = self.wo(x)
        return x


class T5DenseGatedActDense(torch.nn.Module):
    def __init__(self, model_dim, ff_dim, ff_activation):
        super().__init__()
        self.wi_0 = torch.nn.Linear(model_dim, ff_dim, bias=False)
        self.wi_1 = torch.nn.Linear(model_dim, ff_dim, bias=False)
        self.wo = torch.nn.Linear(ff_dim, model_dim, bias=False)
        self.act = activations[ff_activation]

    def forward(self, x):
        hidden_gelu = self.act(self.wi_0(x))
        hidden_linear = self.wi_1(x)
        x = hidden_gelu * hidden_linear
        x = self.wo(x)
        return x


class T5LayerFF(torch.nn.Module):
    def __init__(self, model_dim, ff_dim, ff_activation, gated_act):
        super().__init__()
        if gated_act:
            self.DenseReluDense = T5DenseGatedActDense(model_dim, ff_dim, ff_activation)
        else:
            self.DenseReluDense = T5DenseActDense(model_dim, ff_dim, ff_activation)

        self.layer_norm = T5LayerNorm(model_dim)

    def forward(self, x):
        forwarded_states = self.layer_norm(x)
        forwarded_states = self.DenseReluDense(forwarded_states)
        x += forwarded_states
        return x


class T5Attention(torch.nn.Module):
    def __init__(self, model_dim, inner_dim, num_heads, relative_attention_bias):
        super().__init__()
        self.q = torch.nn.Linear(model_dim, inner_dim, bias=False)
        self.k = torch.nn.Linear(model_dim, inner_dim, bias=False)
        self.v = torch.nn.Linear(model_dim, inner_dim, bias=False)
        self.o = torch.nn.Linear(inner_dim, model_dim, bias=False)
        self.num_heads = num_heads

        self.relative_attention_bias = None
        if relative_attention_bias:
            self.relative_attention_num_buckets = 32
            self.relative_attention_max_distance = 128
            self.relative_attention_bias = torch.nn.Embedding(self.relative_attention_num_buckets, self.num_heads)

    @staticmethod
    def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
        relative_buckets = 0
        if bidirectional:
            num_buckets //= 2
            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
            relative_position = torch.abs(relative_position)
        else:
            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))

        max_exact = num_buckets // 2
        is_small = relative_position < max_exact

        relative_position_if_large = max_exact + (
                torch.log(relative_position.float() / max_exact)
                / math.log(max_distance / max_exact)
                * (num_buckets - max_exact)
        ).to(torch.long)
        relative_position_if_large = torch.min(
            relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)
        )

        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
        return relative_buckets

    def compute_bias(self, query_length, key_length, device, dtype):
        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
        relative_position = memory_position - context_position
        relative_position_bucket = self._relative_position_bucket(
            relative_position,
            bidirectional=True,
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )
        values = self.relative_attention_bias(relative_position_bucket).to(dtype)
        values = values.permute([2, 0, 1]).unsqueeze(0)
        return values

    def forward(self, x, mask=None, past_bias=None):
        q = self.q(x)
        k = self.k(x)
        v = self.v(x)
        if self.relative_attention_bias is not None:
            past_bias = self.compute_bias(x.shape[1], x.shape[1], x.device, x.dtype)

        if past_bias is not None:
            if mask is not None:
                mask = mask + past_bias
            else:
                mask = past_bias

        out = attention_function(q, k * ((k.shape[-1] / self.num_heads) ** 0.5), v, self.num_heads, mask)
        return self.o(out), past_bias


class T5LayerSelfAttention(torch.nn.Module):
    def __init__(self, model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias):
        super().__init__()
        self.SelfAttention = T5Attention(model_dim, inner_dim, num_heads, relative_attention_bias)
        self.layer_norm = T5LayerNorm(model_dim)

    def forward(self, x, mask=None, past_bias=None):
        output, past_bias = self.SelfAttention(self.layer_norm(x), mask=mask, past_bias=past_bias)
        x += output
        return x, past_bias


class T5Block(torch.nn.Module):
    def __init__(self, model_dim, inner_dim, ff_dim, ff_activation, gated_act, num_heads, relative_attention_bias):
        super().__init__()
        self.layer = torch.nn.ModuleList()
        self.layer.append(T5LayerSelfAttention(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias))
        self.layer.append(T5LayerFF(model_dim, ff_dim, ff_activation, gated_act))

    def forward(self, x, mask=None, past_bias=None):
        x, past_bias = self.layer[0](x, mask, past_bias)
        x = self.layer[-1](x)
        return x, past_bias


class T5Stack(torch.nn.Module):
    def __init__(self, num_layers, model_dim, inner_dim, ff_dim, ff_activation, gated_act, num_heads, relative_attention):
        super().__init__()

        self.block = torch.nn.ModuleList(
            [T5Block(model_dim, inner_dim, ff_dim, ff_activation, gated_act, num_heads, relative_attention_bias=((not relative_attention) or (i == 0))) for i in range(num_layers)]
        )
        self.final_layer_norm = T5LayerNorm(model_dim)

    def forward(self, x, attention_mask=None):
        mask = None

        if attention_mask is not None:
            mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])
            mask = mask.masked_fill(mask.to(torch.bool), float("-inf"))

        past_bias = None

        for i, l in enumerate(self.block):
            x, past_bias = l(x, mask, past_bias)

        x = self.final_layer_norm(x)
        return x


class T5(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_layers = config["num_layers"]
        model_dim = config["d_model"]

        self.encoder = T5Stack(self.num_layers, model_dim, model_dim, config["d_ff"], config["dense_act_fn"], config["is_gated_act"], config["num_heads"], config["model_type"] != "umt5")
        self.shared = torch.nn.Embedding(config["vocab_size"], model_dim)

    def forward(self, input_ids, *args, **kwargs):
        x = self.shared(input_ids)
        x = torch.nan_to_num(x)
        return self.encoder(x, *args, **kwargs)


class IntegratedT5(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.transformer = T5(config)
        self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))