backends/python/server/text_embeddings_server/models/jinaBert_model.py

import torch
import math
from torch import nn
import torch.nn.functional as F
from pathlib import Path
from typing import Type, List, Optional, Union, Tuple
from transformers import AutoConfig, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
from opentelemetry import trace
from safetensors import safe_open
from text_embeddings_server.models.pooling import DefaultPooling

from text_embeddings_server.models import Model
from text_embeddings_server.models.types import PaddedBatch, Embedding, Score

tracer = trace.get_tracer(__name__)


class JinaBertConfig(PretrainedConfig):
    def __init__(
        self,
        vocab_size=30522,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_act="gelu",
        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,
        position_embedding_type="absolute",
        use_cache=True,
        classifier_dropout=None,
        feed_forward_type="original",
        emb_pooler=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.position_embedding_type = position_embedding_type
        self.use_cache = use_cache
        self.classifier_dropout = classifier_dropout
        self.feed_forward_type = feed_forward_type
        self.emb_pooler = emb_pooler


class JinaBertEmbeddings:
    """Construct the embeddings from word, position and token_type embeddings."""

    def __init__(self, handle, device, dtype, config: JinaBertConfig):
        self.word_embeddings_weight = (
            handle.get_tensor(f"embeddings.word_embeddings.weight").to(dtype).to(device)
        )
        self.token_type_embeddings_weight = (
            handle.get_tensor(f"embeddings.token_type_embeddings.weight")
            .to(dtype)
            .to(device)
        )
        self.layernorm_weight = (
            handle.get_tensor(f"embeddings.LayerNorm.weight").to(dtype).to(device)
        )
        self.layernorm_bias = (
            handle.get_tensor(f"embeddings.LayerNorm.bias").to(dtype).to(device)
        )
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
        self.position_embedding_type = getattr(
            config, "position_embedding_type", "absolute"
        )
        self.config = config

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
    ) -> torch.Tensor:
        inputs_embeds = F.embedding(input_ids, self.word_embeddings_weight)
        token_type_embeddings = F.embedding(
            token_type_ids, self.token_type_embeddings_weight
        )

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = F.layer_norm(
            embeddings,
            self.layernorm_weight.shape,
            self.layernorm_weight,
            self.layernorm_bias,
            eps=self.config.layer_norm_eps,
        )
        embeddings = self.dropout(embeddings)
        return embeddings


class JinaBertSelfAttention:
    def __init__(self, prefix, handle, device, dtype, config: JinaBertConfig):
        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.config = config
        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_weight = (
            handle.get_tensor(f"{prefix}.query.weight").to(dtype).to(device)
        )
        self.query_bias = handle.get_tensor(f"{prefix}.query.bias").to(dtype).to(device)
        self.key_weight = handle.get_tensor(f"{prefix}.key.weight").to(dtype).to(device)
        self.key_bias = handle.get_tensor(f"{prefix}.key.bias").to(dtype).to(device)
        self.value_weight = (
            handle.get_tensor(f"{prefix}.value.weight").to(dtype).to(device)
        )
        self.value_bias = handle.get_tensor(f"{prefix}.value.bias").to(dtype).to(device)
        self.layer_norm_q_weight = (
            handle.get_tensor(f"{prefix}.layer_norm_q.weight").to(dtype).to(device)
        )
        self.layer_norm_q_bias = (
            handle.get_tensor(f"{prefix}.layer_norm_q.bias").to(dtype).to(device)
        )
        self.layer_norm_k_weight = (
            handle.get_tensor(f"{prefix}.layer_norm_k.weight").to(dtype).to(device)
        )
        self.layer_norm_k_bias = (
            handle.get_tensor(f"{prefix}.layer_norm_k.bias").to(dtype).to(device)
        )

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

    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,
        bias: Optional[torch.FloatTensor] = None,
    ) -> Tuple[torch.Tensor]:
        q_hidden_states = F.linear(hidden_states, self.query_weight, self.query_bias)
        mixed_query_layer = F.layer_norm(
            q_hidden_states,
            self.layer_norm_q_weight.shape,
            self.layer_norm_q_weight,
            self.layer_norm_q_bias,
            eps=self.config.layer_norm_eps,
        )

        k_hidden_states = F.linear(hidden_states, self.key_weight, self.key_bias)
        key_layer = self.transpose_for_scores(
            F.layer_norm(
                k_hidden_states,
                self.layer_norm_k_weight.shape,
                self.layer_norm_k_weight,
                self.layer_norm_k_bias,
                eps=self.config.layer_norm_eps,
            )
        )

        v_hidden_states = F.linear(hidden_states, self.value_weight, self.value_bias)
        value_layer = self.transpose_for_scores(v_hidden_states)
        query_layer = self.transpose_for_scores(mixed_query_layer)

        # 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 = F.softmax(attention_scores + bias, dim=-1)

        attention_probs = self.dropout(attention_probs)

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

        return outputs


class JinaBertSelfOutput:
    def __init__(self, prefix, handle, device, dtype, config):
        self.config = config
        self.dense_weight = (
            handle.get_tensor(f"{prefix}.dense.weight").to(dtype).to(device)
        )
        self.dense_bias = handle.get_tensor(f"{prefix}.dense.bias").to(dtype).to(device)
        self.layerNorm_weight = (
            handle.get_tensor(f"{prefix}.LayerNorm.weight").to(dtype).to(device)
        )
        self.layerNorm_bias = (
            handle.get_tensor(f"{prefix}.LayerNorm.bias").to(dtype).to(device)
        )

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

    def forward(
        self, hidden_states: torch.Tensor, input_tensor: torch.Tensor
    ) -> torch.Tensor:
        hidden_states = F.linear(hidden_states, self.dense_weight, self.dense_bias)
        hidden_states = self.dropout(hidden_states)
        hidden_states = F.layer_norm(
            hidden_states + input_tensor,
            self.layerNorm_weight.shape,
            self.layerNorm_weight,
            self.layerNorm_bias,
            eps=self.config.layer_norm_eps,
        )
        return hidden_states


class JinaBertAttention:
    def __init__(self, prefix, handle, device, dtype, config):
        self.self = JinaBertSelfAttention(
            f"{prefix}.self", handle, device, dtype, config
        )
        self.output = JinaBertSelfOutput(
            f"{prefix}.output", handle, device, dtype, config
        )

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


class JinaBertGLUMLP:
    def __init__(self, prefix, handle, device, dtype, config: JinaBertConfig):
        self.config = config
        if config.feed_forward_type == "reglu":
            self.act = nn.ReLU()
        elif config.feed_forward_type == "geglu":
            self.act = nn.GELU()
        else:
            raise ValueError(
                f"feed_forward_type {config.feed_forward_type} not supported"
            )
        self.up_gated_layer_weight = (
            handle.get_tensor(f"{prefix}.up_gated_layer.weight").to(dtype).to(device)
        )
        self.down_layer_weight = (
            handle.get_tensor(f"{prefix}.down_layer.weight").to(dtype).to(device)
        )
        self.down_layer_bias = (
            handle.get_tensor(f"{prefix}.down_layer.bias").to(dtype).to(device)
        )
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # Up with gate
        hidden_mlp_states = F.linear(hidden_states, self.up_gated_layer_weight, None)
        up = hidden_mlp_states[:, :, : self.config.intermediate_size]
        gated = hidden_mlp_states[:, :, self.config.intermediate_size :]
        hidden_mlp_states = up * self.act(gated)
        hidden_mlp_states = self.dropout(hidden_mlp_states)
        # Down
        return F.linear(hidden_mlp_states, self.down_layer_weight, self.down_layer_bias)


class JinaBertLayer:
    def __init__(self, prefix, handle, device, dtype, config: JinaBertConfig):
        self.attention = JinaBertAttention(
            f"{prefix}.attention", handle, device, dtype, config
        )
        self.config = config
        self.feed_forward_type = config.feed_forward_type
        self.layer_norm_1_weight = (
            handle.get_tensor(f"{prefix}.layer_norm_1.weight").to(dtype).to(device)
        )
        self.layer_norm_1_bias = (
            handle.get_tensor(f"{prefix}.layer_norm_1.bias").to(dtype).to(device)
        )
        self.layer_norm_2_weight = (
            handle.get_tensor(f"{prefix}.layer_norm_2.weight").to(dtype).to(device)
        )
        self.layer_norm_2_bias = (
            handle.get_tensor(f"{prefix}.layer_norm_2.bias").to(dtype).to(device)
        )
        if self.feed_forward_type.endswith("glu"):
            self.mlp = JinaBertGLUMLP(f"{prefix}.mlp", handle, device, dtype, config)
        else:
            raise ValueError(
                f"feed_forward_type {self.feed_forward_type} not supported"
            )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        bias: Optional[torch.FloatTensor] = None,
    ) -> Tuple[torch.Tensor]:
        # Pre-Norm
        residual = hidden_states

        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attention_outputs = self.attention.forward(
            hidden_states,
            attention_mask,
            bias=bias,
        )
        attention_output = self_attention_outputs[0]

        outputs = self_attention_outputs[
            1:
        ]  # add self attentions if we output attention weights

        residual = F.layer_norm(
            residual + attention_output,
            self.layer_norm_1_weight.shape,
            self.layer_norm_1_weight,
            self.layer_norm_1_bias,
            eps=self.config.layer_norm_eps,
        )
        mlp_output = self.mlp.forward(residual)
        layer_output = F.layer_norm(
            residual + mlp_output,
            self.layer_norm_2_weight.shape,
            self.layer_norm_2_weight,
            self.layer_norm_2_bias,
            eps=self.config.layer_norm_eps,
        )
        outputs = (layer_output,) + outputs

        return outputs


class JinaBertEncoder:
    def __init__(self, handle, device, dtype, config: JinaBertConfig):
        self.config = config
        self.layers = [
            JinaBertLayer(f"encoder.layer.{i}", handle, device, dtype, config)
            for i in range(config.num_hidden_layers)
        ]
        self.num_attention_heads = config.num_attention_heads

    def rebuild_alibi_tensor(
        self, size: int, device: Optional[Union[torch.device, str]] = None
    ):
        # Alibi
        # Following https://github.com/ofirpress/attention_with_linear_biases/issues/5 (Implementation 1)
        # In the causal case, you can exploit the fact that softmax is invariant to a uniform translation
        # of the logits, which makes the math work out *after* applying causal masking. If no causal masking
        # will be applied, it is necessary to construct the diagonal mask.
        n_heads = self.num_attention_heads

        def _get_alibi_head_slopes(n_heads: int) -> List[float]:
            def get_slopes_power_of_2(n):
                start = 2 ** (-(2 ** -(math.log2(n) - 3)))
                ratio = start
                return [start * ratio**i for i in range(n)]

            if math.log2(n_heads).is_integer():
                return get_slopes_power_of_2(
                    n_heads
                )  # In the paper, we only train models that have 2^a heads for some a. This function has
            else:  # some good properties that only occur when the input is a power of 2. To maintain that even
                closest_power_of_2 = (
                    2 ** math.floor(math.log2(n_heads))
                )  # when the number of heads is not a power of 2, we use this workaround.
                return (
                    get_slopes_power_of_2(closest_power_of_2)
                    + _get_alibi_head_slopes(2 * closest_power_of_2)[0::2][
                        : n_heads - closest_power_of_2
                    ]
                )

        context_position = torch.arange(size, device=device)[:, None]
        memory_position = torch.arange(size, device=device)[None, :]
        relative_position = torch.abs(memory_position - context_position)
        # [n_heads, max_token_length, max_token_length]
        relative_position = relative_position.unsqueeze(0).expand(n_heads, -1, -1)
        slopes = torch.Tensor(_get_alibi_head_slopes(n_heads)).to(device) * -1
        alibi = slopes.unsqueeze(1).unsqueeze(1) * relative_position
        # [1, n_heads, max_token_length, max_token_length]
        alibi = alibi.unsqueeze(0)
        assert alibi.shape == torch.Size([1, n_heads, size, size])

        self._current_alibi_size = size
        return alibi

    def forward(
        self,
        hidden_states: torch.Tensor,
        max_len: int,
        attention_mask: Optional[torch.FloatTensor] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
        # Add alibi matrix to extended_attention_mask
        bs, seqlen, _ = hidden_states.size()
        alibi_bias = self.rebuild_alibi_tensor(
            size=max_len, device=hidden_states.device
        ).to(hidden_states.dtype)
        full_alibi_bias = torch.full(
            (bs, self.num_attention_heads, seqlen, seqlen),
            fill_value=torch.finfo(hidden_states.dtype).min,
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        full_alibi_bias[:, :, :max_len, :max_len] = alibi_bias

        for i, layer_module in enumerate(self.layers):
            layer_outputs = layer_module.forward(
                hidden_states,
                attention_mask,
                full_alibi_bias,
            )

            hidden_states = layer_outputs[0]

        return hidden_states


class FlashJinaBertModel:
    def __init__(self, handle, device, dtype, config: AutoConfig):
        self.embeddings = JinaBertEmbeddings(handle, device, dtype, config)
        self.encoder = JinaBertEncoder(handle, device, dtype, config)

    def forward(
        self,
        input_ids,
        token_type_ids,
        position_ids,
        max_len,
        attn_mask=None,
    ):
        embeddings = self.embeddings.forward(input_ids, token_type_ids, position_ids)
        encoder_outputs = self.encoder.forward(embeddings, max_len, attn_mask)
        return encoder_outputs


class FlashJinaBert(Model):
    def __init__(
        self,
        model_path: Path,
        device: torch.device,
        dtype: torch.dtype,
        pool: str = "mean",
        trust_remote: bool = True,
    ):
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=trust_remote)
        if hasattr(config, "max_seq_length"):
            self.max_input_length = config.max_seq_length
        else:
            self.max_input_length = config.max_position_embeddings

        with safe_open(model_path / "model.safetensors", framework="pt") as f:
            model = FlashJinaBertModel(f, device, dtype, config)
        self.hidden_size = config.hidden_size
        self.pooling = DefaultPooling(self.hidden_size, pooling_mode=pool)
        self.device = device
        self.dtype = dtype
        self.hidden_size = config.hidden_size

        super(FlashJinaBert, self).__init__(model=model, dtype=dtype, device=device)

    @property
    def batch_type(self) -> Type[PaddedBatch]:
        return PaddedBatch

    def mean_pooling(
        self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor
    ):
        input_mask_expanded = (
            attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        )
        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
            input_mask_expanded.sum(1), min=1e-9
        )

    @tracer.start_as_current_span("embed")
    def embed(self, batch: PaddedBatch) -> List[Embedding]:
        kwargs = {"input_ids": batch.input_ids}
        kwargs["token_type_ids"] = batch.token_type_ids
        kwargs["position_ids"] = batch.position_ids
        input_lens = batch.attention_mask.cumsum(-1)[:, -1].to(torch.int32)
        max_input_lens = input_lens.max().item()
        kwargs["max_len"] = max_input_lens
        outputs = self.model.forward(**kwargs)

        embedding = self.mean_pooling(outputs, batch.attention_mask)
        cpu_results = embedding.view(-1).tolist()

        return [
            Embedding(
                values=cpu_results[i * self.hidden_size : (i + 1) * self.hidden_size]
            )
            for i in range(len(batch))
        ]

    @tracer.start_as_current_span("predict")
    def predict(self, batch: PaddedBatch) -> List[Score]:
        pass