Pool [CLS] token from DistilBERT
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 embeddingI’ve taken this idea from the following DPR code:github.comhuggingface/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,
2022-01-18T20:34:32Z
did you get a solution for your question, please