| """HuggingFace Transformers config for Ivme-Conversate-v2.""" | |
| from transformers import PretrainedConfig | |
| class IvmeConfig(PretrainedConfig): | |
| model_type = "ivme" | |
| def __init__( | |
| self, | |
| vocab_size: int = 16_000, | |
| hidden_dim: int = 384, | |
| n_layers: int = 10, | |
| n_heads: int = 6, | |
| context_len: int = 1024, | |
| ffn_mult: float = 4.0, | |
| rope_theta: float = 10_000.0, | |
| norm_eps: float = 1e-5, | |
| tie_embeddings: bool = True, | |
| dropout: float = 0.0, | |
| **kwargs, | |
| ): | |
| self.vocab_size = vocab_size | |
| self.hidden_dim = hidden_dim | |
| self.n_layers = n_layers | |
| self.n_heads = n_heads | |
| self.context_len = context_len | |
| self.ffn_mult = ffn_mult | |
| self.rope_theta = rope_theta | |
| self.norm_eps = norm_eps | |
| self.dropout = dropout | |
| assert hidden_dim % n_heads == 0, "hidden_dim must be divisible by n_heads" | |
| self.max_position_embeddings = context_len | |
| self.num_hidden_layers = n_layers | |
| self.num_attention_heads = n_heads | |
| self.hidden_size = hidden_dim | |
| kwargs.setdefault("tie_word_embeddings", tie_embeddings) | |
| super().__init__(**kwargs) | |
| def head_dim(self) -> int: | |
| return self.hidden_dim // self.n_heads | |