LightNovel-Intro-RetNet-400M / configuration_retnet.py

Upload 2 files

4371fee verified 4 months ago

No virus

5.03 kB

	# by syncdoth: https://github.com/syncdoth/RetNet/blob/main/retnet/configuration_retnet.py

	from dataclasses import dataclass
	import json

	from transformers.configuration_utils import PretrainedConfig


	def load_config_from_json(config_file):
	with open(config_file, "r") as f:
	config = json.load(f)
	config = RetNetConfig.from_dict(config)
	return config


	@dataclass
	class RetNetConfig(PretrainedConfig):
	model_type = "retnet"
	initializer_range: float = 0.02
	activation_fn: str = "gelu"
	dropout: float = 0.0 # dropout probability
	activation_dropout: float = 0.0 # dropout probability after activation in FFN.
	drop_path_rate: float = 0.0
	decoder_embed_dim: int = 768 # decoder embedding dimension
	decoder_value_embed_dim: int = 1280 # decoder value embedding dimension
	decoder_ffn_embed_dim: int = 1280 # decoder embedding dimension for FFN
	decoder_layers: int = 12 # num decoder layers
	decoder_retention_heads: int = 3 # num decoder retention heads
	decoder_normalize_before: bool = True # apply layernorm before each decoder block
	layernorm_embedding: bool = False # add layernorm to embedding
	no_scale_embedding: bool = True # if True, dont scale embeddings
	recurrent_chunk_size: int = 512
	use_lm_decay: bool = False
	use_glu: bool = True # use GLU instead of FFN
	z_loss_coeff: float = 0.0 # coefficient for z loss: TODO: 1e-4
	deepnorm: bool = False
	subln: bool = True
	use_ffn_rms_norm: bool = False
	layernorm_eps: float = 1e-6
	tie_word_embeddings: bool = False

	def __init__(
	self,
	vocab_size: int = 50257,
	initializer_range: float = 0.02,
	is_decoder: bool = True,
	pad_token_id: int = 0,
	eos_token_id: int = 0,
	output_retentions: bool = False,
	use_cache: bool = True,
	forward_impl: str = "parallel",
	activation_fn: str = "gelu",
	dropout: float = 0.0, # dropout probability
	activation_dropout: float = 0.0, # dropout probability after activation in FFN.
	drop_path_rate: float = 0.0,
	decoder_embed_dim: int = 768, # decoder embedding dimension
	decoder_value_embed_dim: int = 1280, # decoder value embedding dimension
	decoder_ffn_embed_dim: int = 1280, # decoder embedding dimension for FFN
	decoder_layers: int = 12, # num decoder layers
	decoder_retention_heads: int = 3, # num decoder retention heads
	decoder_normalize_before: bool = True, # apply layernorm before each decoder block
	layernorm_embedding: bool = False, # add layernorm to embedding
	no_scale_embedding: bool = True, # if True, dont scale embeddings
	recurrent_chunk_size: int = 512,
	use_glu: bool = True, # use GLU instead of FFN
	z_loss_coeff: float = 0.0, # coefficient for z loss: TODO: 1e-4
	use_lm_decay: bool = False,
	deepnorm: bool = False,
	subln: bool = True,
	use_ffn_rms_norm: bool = False, # use RMSNorm instead of LayerNorm in FFN
	layernorm_eps: float = 1e-6,
	tie_word_embeddings: bool = False,
	**kwargs
	):
	self.vocab_size = vocab_size
	self.initializer_range = initializer_range
	self.output_retentions = output_retentions
	self.use_lm_decay = use_lm_decay
	self.use_glu = use_glu
	self.z_loss_coeff = z_loss_coeff
	# size related
	self.decoder_embed_dim = decoder_embed_dim
	self.decoder_value_embed_dim = decoder_value_embed_dim
	self.decoder_retention_heads = decoder_retention_heads
	self.decoder_ffn_embed_dim = decoder_ffn_embed_dim
	self.decoder_layers = decoder_layers
	# normalization related
	self.decoder_normalize_before = decoder_normalize_before
	self.activation_fn = activation_fn
	self.dropout = dropout
	self.drop_path_rate = drop_path_rate
	self.activation_dropout = activation_dropout
	self.no_scale_embedding = no_scale_embedding
	self.layernorm_embedding = layernorm_embedding
	self.deepnorm = deepnorm
	self.subln = subln
	self.use_ffn_rms_norm = use_ffn_rms_norm
	self.layernorm_eps = layernorm_eps
	# Blockwise
	self.recurrent_chunk_size = recurrent_chunk_size
	self.forward_impl = forward_impl

	if self.deepnorm:
	self.decoder_normalize_before = False
	self.subln = False
	if self.subln:
	self.decoder_normalize_before = True
	self.deepnorm = False

	super().__init__(
	is_decoder=is_decoder,
	pad_token_id=pad_token_id,
	eos_token_id=eos_token_id,
	use_cache=use_cache,
	tie_word_embeddings=tie_word_embeddings,
	**kwargs
	)

	def override(self, args):
	for hp in self.__dict__.keys():
	if getattr(args, hp, None) is not None:
	self.__dict__[hp] = getattr(args, hp, None)