Pool [CLS] token from DistilBERT

AndreGodinho · October 31, 2021, 11:29pm

Hey there everyone, I haven’t used this amazing API for a while. The code below is able to pool the [CLS] output embedding from DistilBERT?

distilbert = DistilBertModel.from_pretrained('distilbert-base-uncased', return_dict=True)

# more code ...

# input_ids & attention_mask from a batch
outputs = distilbert(input_ids, attention_mask=attention_mask)
seq_embeddings = outputs[0] # DistilBERT output embeddings 
cls_embedding = seq_embeddings[:,0,:] # <--- [CLS] output embedding

I’ve taken this idea from the following DPR code:

github.com

huggingface/transformers/blob/master/src/transformers/models/dpr/modeling_dpr.py#L204

    
      
          ) -> Union[BaseModelOutputWithPooling, Tuple[Tensor, ...]]:
              outputs = self.bert_model(
                  input_ids=input_ids,
                  attention_mask=attention_mask,
                  token_type_ids=token_type_ids,
                  inputs_embeds=inputs_embeds,
                  output_attentions=output_attentions,
                  output_hidden_states=output_hidden_states,
                  return_dict=return_dict,
              )
              sequence_output, pooled_output = outputs[:2]
              pooled_output = sequence_output[:, 0, :]
              if self.projection_dim > 0:
                  pooled_output = self.encode_proj(pooled_output)
          
          
    if not return_dict:
                  return (sequence_output, pooled_output) + outputs[2:]
          
          
    return BaseModelOutputWithPooling(
                  last_hidden_state=sequence_output,
                  pooler_output=pooled_output,

toyl · January 18, 2022, 8:34pm

did you get a solution for your question, please

Topic		Replies	Views
DistilBERT and CLS token Beginners	2	1997	February 21, 2021
How to properly compute Sentence Embeddings using a non english, pretrained distilbert model? Beginners	0	457	April 25, 2021
Why is using my DistilBERT model for inference so slow? Intermediate	0	816	June 18, 2021
Accessing uncontextualized BERT word embeddings Beginners	2	1412	October 30, 2020
Extracting embeddings with distilbert? (in tensorflow) 🤗Transformers	5	2048	August 6, 2021

Pool [CLS] token from DistilBERT

Related Topics