|
|
from transformers import PretrainedConfig
|
|
|
import json
|
|
|
|
|
|
|
|
|
class HyenaConfig(PretrainedConfig):
|
|
|
model_type = "hyenadna"
|
|
|
def __init__(
|
|
|
self,
|
|
|
vocab_size=12,
|
|
|
d_model=256,
|
|
|
d_inner=None,
|
|
|
use_bias=True,
|
|
|
train_freq=True,
|
|
|
max_seq_len=1024,
|
|
|
emb_dim=3,
|
|
|
n_layer=12,
|
|
|
num_inner_mlps=2,
|
|
|
hyena_order=2,
|
|
|
short_filter_order=3,
|
|
|
filter_order=64,
|
|
|
activation_freq=1,
|
|
|
embed_dropout=0.1,
|
|
|
hyena_dropout=0.0,
|
|
|
hyena_filter_dropout=0.0,
|
|
|
layer_norm_epsilon=1e-5,
|
|
|
initializer_range=0.02,
|
|
|
pad_vocab_size_multiple=8,
|
|
|
**kwargs,
|
|
|
):
|
|
|
self.vocab_size = vocab_size
|
|
|
self.d_model = d_model
|
|
|
if d_inner is None:
|
|
|
self.d_inner = 4 * d_model
|
|
|
else:
|
|
|
self.d_inner = d_inner
|
|
|
self.use_bias = use_bias
|
|
|
self.train_freq = train_freq
|
|
|
self.max_seq_len = max_seq_len
|
|
|
self.emb_dim = emb_dim
|
|
|
self.n_layer = n_layer
|
|
|
self.hyena_order = hyena_order
|
|
|
self.filter_order = filter_order
|
|
|
self.short_filter_order = short_filter_order
|
|
|
self.activation_freq = activation_freq
|
|
|
self.num_inner_mlps = num_inner_mlps
|
|
|
self.embed_dropout = embed_dropout
|
|
|
self.hyena_dropout = hyena_dropout
|
|
|
self.hyena_filter_dropout = hyena_filter_dropout
|
|
|
self.layer_norm_epsilon = layer_norm_epsilon
|
|
|
self.initializer_range = initializer_range
|
|
|
self.pad_vocab_size_multiple = pad_vocab_size_multiple
|
|
|
super().__init__(**kwargs)
|
|
|
|
|
|
@classmethod
|
|
|
def from_original_config(cls, config_path, **kwargs):
|
|
|
with open(config_path, "r") as f:
|
|
|
config = json.load(f)
|
|
|
|
|
|
vocab_size = config["vocab_size"]
|
|
|
d_model = config["d_model"]
|
|
|
d_inner = config["d_inner"]
|
|
|
max_seq_len = config["layer"]["l_max"]
|
|
|
emb_dim = config["layer"]["emb_dim"]
|
|
|
filter_order = config["layer"]["filter_order"]
|
|
|
if "local_order" in config["layer"]:
|
|
|
short_filter_order = config["layer"]["local_order"]
|
|
|
elif "short_filter_order" in config["layer"]:
|
|
|
short_filter_order = config["layer"]["short_filter_order"]
|
|
|
else:
|
|
|
short_filter_order = 3
|
|
|
n_layer = config["n_layer"]
|
|
|
activation_freq = config["layer"]["w"]
|
|
|
embed_dropout = config["embed_dropout"]
|
|
|
pad_vocab_size_multiple = config["pad_vocab_size_multiple"]
|
|
|
return cls(vocab_size=vocab_size,
|
|
|
d_model=d_model,
|
|
|
d_inner=d_inner,
|
|
|
max_seq_len=max_seq_len,
|
|
|
emb_dim=emb_dim,
|
|
|
filter_order=filter_order,
|
|
|
short_filter_order=short_filter_order,
|
|
|
n_layer=n_layer,
|
|
|
activation_freq=activation_freq,
|
|
|
embed_dropout=embed_dropout,
|
|
|
pad_vocab_size_multiple=pad_vocab_size_multiple,
|
|
|
tie_word_embeddings=False,
|
|
|
**kwargs
|
|
|
) |
