vortex_config.py

"""
Vortex-7B model configuration.
Optimized for 8GB VRAM (4060 laptop) and MacBook Pro M2/M3.
"""

VORTEX_7B_CONFIG = {
    # Model dimensions
    "d_model": 4096,
    "num_layers": 32,
    "num_heads": 32,
    "head_dim": 128,  # d_model // num_heads

    # State-space layer parameters
    "d_state": 16,          # SSM state dimension
    "d_conv": 4,            # SSM convolution width

    # Attention parameters
    "window_size": 512,     # Local attention window
    "use_flash_attention": True,  # CUDA only

    # Feed-forward parameters
    "ffn_expansion": 4,     # Hidden dim = d_model * expansion
    "num_domains": 7,       # Physics, Math, Chemistry, Biology, Earth, Space, Zoology

    # Tokenizer parameters
    "vocab_size": 50000,
    "max_seq_len": 16384,

    # Layer ratio: 60% SSM, 40% attention
    "ssm_ratio": 0.6,

    # Data types
    "dtype": "bfloat16",

    # Special tokens
    "special_tokens": {
        "[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,
    },

    # Domain tags
    "domain_tags": ["[MATH]", "[CHEM]", "[BIO]", "[PHYS]", "[EARTH]", "[SPACE]", "[ZOO]"],

    # Science module flags (enable/disable for ablation)
    "enable_equation_module": True,
    "enable_numerical_module": True,
    "enable_citation_module": True,
    "enable_molecular_module": True,
}


def get_config():
    """Return the 7B configuration dictionary."""
    return VORTEX_7B_CONFIG