| | """
|
| | Vortex configuration for HuggingFace.
|
| | """
|
| |
|
| | from typing import Optional, List, Dict, Any
|
| | from transformers import PretrainedConfig
|
| |
|
| |
|
| | class VortexConfig(PretrainedConfig):
|
| | """
|
| | Configuration class for Vortex model.
|
| | Compatible with HuggingFace transformers.
|
| | """
|
| |
|
| | model_type = "vortex"
|
| | tie_word_embeddings = True
|
| |
|
| | def __init__(
|
| | self,
|
| | d_model: int = 4096,
|
| | num_layers: int = 32,
|
| | num_heads: int = 32,
|
| | d_state: int = 16,
|
| | d_conv: int = 4,
|
| | window_size: int = 512,
|
| | ffn_expansion: int = 4,
|
| | num_domains: int = 7,
|
| | vocab_size: int = 50000,
|
| | max_seq_len: int = 16384,
|
| | ssm_ratio: float = 0.6,
|
| | enable_equation_module: bool = True,
|
| | enable_numerical_module: bool = True,
|
| | enable_citation_module: bool = True,
|
| | enable_molecular_module: bool = True,
|
| | special_tokens: Optional[Dict[str, int]] = None,
|
| | domain_tags: Optional[List[str]] = None,
|
| | initializer_range: float = 0.02,
|
| | tie_word_embeddings: bool = True,
|
| | **kwargs
|
| | ):
|
| | super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
|
| | self.d_model = d_model
|
| | self.num_layers = num_layers
|
| | self.num_heads = num_heads
|
| | self.d_state = d_state
|
| | self.d_conv = d_conv
|
| | self.window_size = window_size
|
| | self.ffn_expansion = ffn_expansion
|
| | self.num_domains = num_domains
|
| | self.vocab_size = vocab_size
|
| | self.max_seq_len = max_seq_len
|
| | self.ssm_ratio = ssm_ratio
|
| | self.enable_equation_module = enable_equation_module
|
| | self.enable_numerical_module = enable_numerical_module
|
| | self.enable_citation_module = enable_citation_module
|
| | self.enable_molecular_module = enable_molecular_module
|
| | self.special_tokens = special_tokens or {
|
| | "[PAD]": 0, "[UNK]": 1, "[BOS]": 2, "[EOS]": 3,
|
| | "[EQUATION]": 4, "[/EQUATION]": 5,
|
| | "[CITATION]": 6, "[/CITATION]": 7,
|
| | "[MOLECULE]": 8, "[/MOLECULE]": 9,
|
| | "[FIGURE]": 10, "[TABLE]": 11,
|
| | "[MATH]": 12, "[CHEM]": 13, "[BIO]": 14,
|
| | "[PHYS]": 15, "[EARTH]": 16, "[SPACE]": 17, "[ZOO]": 18,
|
| | }
|
| | self.domain_tags = domain_tags or ["[MATH]", "[CHEM]", "[BIO]", "[PHYS]", "[EARTH]", "[SPACE]", "[ZOO]"]
|
| | self.initializer_range = initializer_range
|
| |
|
| | self.head_dim = self.d_model // self.num_heads
|
| |
|
| | @classmethod
|
| | def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
|
| | """Load config from pretrained model."""
|
| | import json
|
| | import os
|
| |
|
| | config_path = os.path.join(pretrained_model_name_or_path, "config.json")
|
| | if os.path.exists(config_path):
|
| | with open(config_path, "r") as f:
|
| | config_dict = json.load(f)
|
| | config_dict.update(kwargs)
|
| | return cls(**config_dict)
|
| | else:
|
| |
|
| | return cls(**kwargs)
|
| |
|
| | def to_dict(self) -> Dict[str, Any]:
|
| | """Convert to dictionary."""
|
| | return {
|
| | "model_type": self.model_type,
|
| | "d_model": self.d_model,
|
| | "num_layers": self.num_layers,
|
| | "num_heads": self.num_heads,
|
| | "head_dim": self.head_dim,
|
| | "d_state": self.d_state,
|
| | "d_conv": self.d_conv,
|
| | "window_size": self.window_size,
|
| | "ffn_expansion": self.ffn_expansion,
|
| | "num_domains": self.num_domains,
|
| | "vocab_size": self.vocab_size,
|
| | "max_seq_len": self.max_seq_len,
|
| | "ssm_ratio": self.ssm_ratio,
|
| | "enable_equation_module": self.enable_equation_module,
|
| | "enable_numerical_module": self.enable_numerical_module,
|
| | "enable_citation_module": self.enable_citation_module,
|
| | "enable_molecular_module": self.enable_molecular_module,
|
| | "special_tokens": self.special_tokens,
|
| | "domain_tags": self.domain_tags,
|
| | "initializer_range": self.initializer_range,
|
| | }
|
| |
|
