| | |
| | |
| | """ PyTorch CodeT5+ mbedding models. |
| | The implementation is based on transformers.models.t5.modeling_t5 by adding a projection layer on T5EncoderModel |
| | """ |
| |
|
| | from typing import Optional, Tuple, Union |
| | import torch |
| | from torch import nn |
| | import torch.nn.functional as F |
| | from transformers import T5EncoderModel |
| | from transformers.modeling_outputs import ( |
| | BaseModelOutput, |
| | ) |
| | from .configuration_codet5p_embedding import CodeT5pEmbeddingConfig |
| |
|
| |
|
| | class CodeT5pEmbeddingModel(T5EncoderModel): |
| | config_class = CodeT5pEmbeddingConfig |
| |
|
| | authorized_missing_keys = [ |
| | r"encoder.embed_tokens.weight", |
| | ] |
| |
|
| | def __init__(self, config: CodeT5pEmbeddingConfig): |
| | super().__init__(config) |
| | self.proj = nn.Linear(config.d_model, config.embed_dim) |
| |
|
| | def forward( |
| | self, |
| | input_ids: Optional[torch.LongTensor] = None, |
| | attention_mask: Optional[torch.FloatTensor] = None, |
| | head_mask: Optional[torch.FloatTensor] = None, |
| | inputs_embeds: Optional[torch.FloatTensor] = None, |
| | output_attentions: Optional[bool] = None, |
| | output_hidden_states: Optional[bool] = None, |
| | return_dict: Optional[bool] = None, |
| | ) -> Union[Tuple[torch.FloatTensor], BaseModelOutput]: |
| | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| |
|
| | encoder_outputs = self.encoder( |
| | input_ids=input_ids, |
| | attention_mask=attention_mask, |
| | inputs_embeds=inputs_embeds, |
| | head_mask=head_mask, |
| | output_attentions=output_attentions, |
| | output_hidden_states=output_hidden_states, |
| | return_dict=return_dict, |
| | ) |
| |
|
| | embedding = F.normalize(self.proj(encoder_outputs.last_hidden_state[:, 0, :]), dim=-1) |
| |
|
| | return embedding |
| |
|
