from transformers import PretrainedConfig from typing import List class LMConfig(PretrainedConfig): model_type = "minimind" def __init__( self, dim: int = 512, n_layers: int = 8, n_heads: int = 8, n_kv_heads: int = 2, vocab_size: int = 25, # 6400 hidden_dim: int = None, logit_dim : int = 9, multiple_of: int = 64, norm_eps: float = 1e-5, max_seq_len: int = 8192, rope_theta: int = 1e6, dropout: float = 0.0, flash_attn: bool = True, #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid #################################################### use_moe: bool = False, #################################################### num_experts_per_tok: int = 2, n_routed_experts: int = 4, n_shared_experts: bool = True, scoring_func: str = 'softmax', aux_loss_alpha: float = 0.1, seq_aux: bool = True, norm_topk_prob: bool = True, padding_idx:int = None, **kwargs, ): self.dim = dim self.n_layers = n_layers self.n_heads = n_heads self.n_kv_heads = n_kv_heads self.vocab_size = vocab_size self.hidden_dim = hidden_dim self.logit_dim = logit_dim self.multiple_of = multiple_of self.norm_eps = norm_eps self.max_seq_len = max_seq_len self.rope_theta = rope_theta self.dropout = dropout self.flash_attn = flash_attn #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid #################################################### self.use_moe = use_moe self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量 self.n_routed_experts = n_routed_experts # 总的专家数量 self.n_shared_experts = n_shared_experts # 共享专家 self.scoring_func = scoring_func # 评分函数，默认为'softmax' self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数 self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失 self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率 #################################################### # Here are the specific configurations of input layer #################################################### self.padding_idx = padding_idx super().__init__(**kwargs) class LMaoTaoConfig(PretrainedConfig): model_type = "minimind" def __init__( self, dim: int = 256, n_layers: int = 8, n_heads: int = 8, n_kv_heads: int = 2, # vocab_size: int = 32, # 6400 vocab_size: int = 9, # 6400 hidden_dim: int = None, logit_dim : int = 32, multiple_of: int = 64, norm_eps: float = 1e-5, max_seq_len: int = 8192, rope_theta: int = 1e6, dropout: float = 0.0, flash_attn: bool = True, #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid #################################################### use_moe: bool = False, #################################################### num_experts_per_tok: int = 2, n_routed_experts: int = 4, n_shared_experts: bool = True, scoring_func: str = 'softmax', aux_loss_alpha: float = 0.1, seq_aux: bool = True, norm_topk_prob: bool = True, padding_idx:int = None, #################################################### ### user defined parameters #################################################### aa_vocab_size: int = 22, # 20 AA + * + pad species_size: int = 5, # 5 species truncated_size: int = 5, # full,head,tail,boundary,middle continuous_features_dim = 3, # data['off_start'],data['off_end'],data['full_len'] log(x+1) head_codon_num:int = 100, **kwargs, ): self.dim = dim self.n_layers = n_layers self.n_heads = n_heads self.n_kv_heads = n_kv_heads self.vocab_size = vocab_size self.hidden_dim = hidden_dim self.logit_dim = logit_dim self.multiple_of = multiple_of self.norm_eps = norm_eps self.max_seq_len = max_seq_len self.rope_theta = rope_theta self.dropout = dropout self.flash_attn = flash_attn #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid #################################################### self.use_moe = use_moe self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量 self.n_routed_experts = n_routed_experts # 总的专家数量 self.n_shared_experts = n_shared_experts # 共享专家 self.scoring_func = scoring_func # 评分函数，默认为'softmax' self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数 self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失 self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率 #################################################### # Here are the specific configurations of input layer #################################################### self.padding_idx = padding_idx #################################################### ### user defined parameters #################################################### self.aa_vocab_size = aa_vocab_size self.species_size = species_size self.truncated_size = truncated_size self.continuous_features_dim = continuous_features_dim self.head_codon_num = head_codon_num super().__init__(**kwargs)