File size: 939 Bytes
9e62c4f 18563f0 9e62c4f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
# coding=utf-8
# Copyright 2023 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
""" PyTorch CodeT5+ matching 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 T5ForConditionalGeneration
from transformers.modeling_outputs import (
BaseModelOutput,
)
from .configuration_codet5p_bimodal import CodeT5pBimodalConfig
class CodeT5pBimodalModel(T5ForConditionalGeneration):
config_class = CodeT5pBimodalConfig
authorized_missing_keys = [
r"encoder.embed_tokens.weight",
]
def __init__(self, config: CodeT5pBimodalConfig):
super().__init__(config)
self.proj = nn.Linear(config.d_model, config.embed_dim)
self.itm_head = nn.Linear(config.d_model, 2)
|