|
|
|
|
|
""" 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 |
|
|
