BertPooler with ggml

[GautamiDeshpande]

Hey, thanks for the great work.

I am trying to implement a bert classifier based on bert-cpp. While most of the classifier layer implementation is done, I am stuck at the pooler layer.

For a bert pooler defined as the following

class BertPooler(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) -> torch.Tensor:
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output

from https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py#L654

I have to use hidden_states[:, 0] variable which will then be passed through a dense pooler, tanh and a dense classifier layer.

I have been able to implement pooler, tanh and dense classifier layer, but I am not sure how I can take only the hidden states corresponding to the first token ie, [CLS] token which will be used for the final classification.

I would really appreciate your inputs.
Thanks in advance!

[monatis]

You can do it with ggml_get_rows function. For example, I'm getting the hidden states at te last index in clip.cpp as in the following:

https://github.com/monatis/clip.cpp/blob/7ca25f046ff1b193d670645b216162c04e9ee5d3/clip.cpp#L941-L945

    // get the output of eot token, e.g., last index
    struct ggml_tensor * eot = ggml_new_i32(ctx0, N - 1);
    embeddings = ggml_get_rows(ctx0, embeddings, eot);

So you can do the following in your case:

    // get the output of cls token, e.g., first index
    struct ggml_tensor * cls = ggml_new_i32(ctx0, 0);
    embeddings = ggml_get_rows(ctx0, embeddings, cls);

[GautamiDeshpande]

@monatis Thanks! I'll try it out.