configs/touchgrass_7b_config.py

"""
TouchGrass-7B model configuration.
Based on Qwen3.5-7B-Instruct with music adaptations.
"""

TOUCHGRASS_7B_CONFIG = {
    # Base model
    "base_model": "Qwen/Qwen3.5-7B-Instruct",
    "model_type": "touchgrass",

    # Model dimensions (from Qwen3.5-7B)
    "d_model": 4096,
    "num_layers": 40,
    "num_heads": 32,
    "head_dim": 128,
    "ffn_expansion": 2.67,  # SwiGLU expansion

    # Tokenizer
    "vocab_size": 32000,  # Qwen3.5 vocab + music tokens
    "max_seq_len": 4096,

    # Music modules
    "enable_tab_chord_module": True,
    "enable_music_theory_module": True,
    "enable_ear_training_module": True,
    "enable_eq_adapter": True,
    "enable_songwriting_module": True,

    # EQ adapter settings
    "eq_hidden_dim": 32,
    "eq_loss_weight": 0.1,

    # Music domain tags
    "music_domains": ["[GUITAR]", "[PIANO]", "[DRUMS]", "[VOCALS]", "[THEORY]", "[DJ]"],
    "skill_levels": ["[BEGINNER]", "[INTERMEDIATE]", "[ADVANCED]"],
    "notation_tags": ["[TAB]", "[CHORD]", "[SHEET]", "[LYRICS]", "[PROGRESSION]"],

    # Special tokens
    "special_tokens": {
        "[PAD]": 0,
        "[UNK]": 1,
        "[BOS]": 2,
        "[EOS]": 3,
        # Music domain tokens
        "[GUITAR]": 32000,
        "[PIANO]": 32001,
        "[DRUMS]": 32002,
        "[VOCALS]": 32003,
        "[THEORY]": 32004,
        "[DJ]": 32005,
        # Notation tokens
        "[TAB]": 32006,
        "[/TAB]": 32007,
        "[CHORD]": 32008,
        "[/CHORD]": 32009,
        "[SHEET]": 32010,
        "[/SHEET]": 32011,
        "[LYRICS]": 32012,
        "[/LYRICS]": 32013,
        "[PROGRESSION]": 32014,
        "[/PROGRESSION]": 32015,
        # Skill level tokens
        "[BEGINNER]": 32016,
        "[INTERMEDIATE]": 32017,
        "[ADVANCED]": 32018,
        # EQ tokens
        "[FRUSTRATED]": 32019,
        "[ENCOURAGED]": 32020,
    },

    # Data types
    "dtype": "bfloat16",

    # Initialization
    "initializer_range": 0.02,
}


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