File size: 2,869 Bytes
1828fa4 |
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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
""" Prot2Text configuration"""
from transformers.configuration_utils import PretrainedConfig
from transformers import AutoConfig
from transformers.utils import logging
logger = logging.get_logger(__name__)
class Prot2TextConfig(PretrainedConfig):
model_type = "prot2text"
keys_to_ignore_at_inference = ["past_key_values"]
_keys_to_ignore_on_load_missing = [r"transformer"]
def __init__(
self,
cross_esm_graph=True,
decoder_start_token_id=50257,
early_stopping=True,
eos_token_id=50258,
bos_token_id=50257,
esm=True,
esm_model_name="facebook/esm2_t6_8M_UR50D",
gpt_model_name="gpt2",
length_penalty=2.0,
max_new_tokens=256,
no_repeat_ngram_size=3,
pad_token_id=50256,
prot2text_version="1.1",
rgcn=True,
rgc_input_dim=67,
rgcn_n_layers=6,
gpt_config=None,
esm_config=None,
**kwargs,
):
self.cross_esm_graph = cross_esm_graph
self.decoder_start_token_id = decoder_start_token_id
self.early_stopping = early_stopping
self.eos_token_id = eos_token_id
self.esm = esm
self.esm_model_name = esm_model_name
self.gpt_model_name = gpt_model_name
self.length_penalty = length_penalty
self.max_new_tokens = max_new_tokens
self.no_repeat_ngram_size = no_repeat_ngram_size
self.pad_token_id = pad_token_id
self.prot2text_version = prot2text_version
self.rgcn = rgcn
self.rgc_input_dim = rgc_input_dim
self.rgcn_n_layers = rgcn_n_layers
if gpt_config is None:
self.gpt_config = AutoConfig.from_pretrained(gpt_model_name,
_name_or_path= gpt_model_name,
is_encoder_decoder=True,
use_cache=False,
add_cross_attention=True,
bos_token_id=bos_token_id,
decoder_start_token_id=decoder_start_token_id,
eos_token_id=eos_token_id,
max_new_tokens=max_new_tokens,
pad_token_id=50256,
vocab_size=50259,
num_beams=1,
max_length=256,
min_length=1).to_dict()
else:
self.gpt_config = gpt_config
if esm_config is None:
self.esm_config = AutoConfig.from_pretrained(esm_model_name).to_dict()
self.esm_config = esm_config
super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) |