"""Centralized registry of the 5 candidate models for WSP v3.0 (9 Subspaces).""" from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import List, Optional PIPELINE_DIR = Path(r"C:\Users\jhjwi\.gemini\antigravity\scratch\wsp-5model-pipeline") MODELS_DIR = PIPELINE_DIR / "models" @dataclass(frozen=True) class ModelSpec: slug: str display_name: str hf_repo: str gguf_repo: Optional[str] gguf_file: str gguf_arch: str num_layers: int hidden_size: int intermediate_size: int tied_embeddings: bool rmsnorm_convention: str # "1+w" or "w" mlp_activation: str # "swiglu" | "relu2" | "geglu" mid_layers: List[int] late_layers: List[int] embed_key: str lm_head_key: Optional[str] norm_f_key: str layer_norm_tmpl: str post_ffn_norm_tmpl: Optional[str] gate_proj_tmpl: Optional[str] up_proj_tmpl: str down_proj_tmpl: str reasoning_format: Optional[str] enable_thinking_kwargs: dict control_family: str # "qwen35_xml" | "gemma4" local_hf_dir: Optional[Path] = None local_gguf_path: Optional[Path] = None @property def work_dir(self) -> Path: return MODELS_DIR / self.slug @property def gguf_path(self) -> Path: return self.local_gguf_path if self.local_gguf_path else (self.work_dir / self.gguf_file) @property def adapter_dir(self) -> Path: return self.work_dir / "adapters" MODEL_SPECS: List[ModelSpec] = [ ModelSpec( slug="qwen35_2b", display_name="Qwen/Qwen3.5-2B", hf_repo="Qwen/Qwen3.5-2B", gguf_repo="unsloth/Qwen3.5-2B-GGUF", gguf_file="Qwen3.5-2B-Q4_K_M.gguf", gguf_arch="qwen35", num_layers=24, hidden_size=2048, intermediate_size=6144, tied_embeddings=True, rmsnorm_convention="1+w", mlp_activation="swiglu", mid_layers=[3, 7, 11, 15], late_layers=[18, 19, 21, 22, 23], embed_key="model.language_model.embed_tokens.weight", lm_head_key=None, norm_f_key="model.language_model.norm.weight", layer_norm_tmpl="model.language_model.layers.{l}.post_attention_layernorm.weight", post_ffn_norm_tmpl=None, gate_proj_tmpl="model.language_model.layers.{l}.mlp.gate_proj.weight", up_proj_tmpl="model.language_model.layers.{l}.mlp.up_proj.weight", down_proj_tmpl="model.language_model.layers.{l}.mlp.down_proj.weight", reasoning_format="deepseek", enable_thinking_kwargs={}, control_family="qwen35_xml", ), ModelSpec( slug="qwen35_4b", display_name="Qwen/Qwen3.5-4B", hf_repo="Qwen/Qwen3.5-4B", gguf_repo=None, gguf_file="Qwen3.5-4B-Q4_K_M.gguf", gguf_arch="qwen35", num_layers=32, hidden_size=2560, intermediate_size=9216, tied_embeddings=True, rmsnorm_convention="1+w", mlp_activation="swiglu", mid_layers=[7, 11, 15, 19], late_layers=[23, 27, 29, 30, 31], embed_key="model.language_model.embed_tokens.weight", lm_head_key=None, norm_f_key="model.language_model.norm.weight", layer_norm_tmpl="model.language_model.layers.{l}.post_attention_layernorm.weight", post_ffn_norm_tmpl=None, gate_proj_tmpl="model.language_model.layers.{l}.mlp.gate_proj.weight", up_proj_tmpl="model.language_model.layers.{l}.mlp.up_proj.weight", down_proj_tmpl="model.language_model.layers.{l}.mlp.down_proj.weight", reasoning_format="deepseek", enable_thinking_kwargs={}, control_family="qwen35_xml", local_hf_dir=Path(r"C:\Users\jhjwi\.gemini\antigravity\scratch\qwen35-4b-wsp\hf_model"), local_gguf_path=Path(r"C:\Users\jhjwi\.gemini\antigravity\scratch\qwen35-4b-wsp\gguf\Qwen3.5-4B-Q4_K_M.gguf"), ), ModelSpec( slug="nemotron3_4b", display_name="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", hf_repo="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16", gguf_repo="nvidia/NVIDIA-Nemotron-3-Nano-4B-GGUF", gguf_file="NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf", gguf_arch="nemotron_h", num_layers=42, hidden_size=3136, intermediate_size=12544, tied_embeddings=False, rmsnorm_convention="w", mlp_activation="relu2", mid_layers=[13, 18, 22, 25], late_layers=[29, 33, 37, 39, 41], embed_key="backbone.embeddings.weight", lm_head_key="lm_head.weight", norm_f_key="backbone.norm_f.weight", layer_norm_tmpl="backbone.layers.{l}.norm.weight", post_ffn_norm_tmpl=None, gate_proj_tmpl=None, up_proj_tmpl="backbone.layers.{l}.mixer.up_proj.weight", down_proj_tmpl="backbone.layers.{l}.mixer.down_proj.weight", reasoning_format="deepseek", enable_thinking_kwargs={"enable_thinking": True}, control_family="qwen35_xml", ), ModelSpec( slug="qwen35_9b", display_name="Qwen/Qwen3.5-9B", hf_repo="Qwen/Qwen3.5-9B", gguf_repo="unsloth/Qwen3.5-9B-GGUF", gguf_file="Qwen3.5-9B-Q4_K_M.gguf", gguf_arch="qwen35", num_layers=32, hidden_size=4096, intermediate_size=12288, tied_embeddings=False, rmsnorm_convention="1+w", mlp_activation="swiglu", mid_layers=[7, 11, 15, 19], late_layers=[23, 27, 29, 30, 31], embed_key="model.language_model.embed_tokens.weight", lm_head_key="lm_head.weight", norm_f_key="model.language_model.norm.weight", layer_norm_tmpl="model.language_model.layers.{l}.post_attention_layernorm.weight", post_ffn_norm_tmpl=None, gate_proj_tmpl="model.language_model.layers.{l}.mlp.gate_proj.weight", up_proj_tmpl="model.language_model.layers.{l}.mlp.up_proj.weight", down_proj_tmpl="model.language_model.layers.{l}.mlp.down_proj.weight", reasoning_format="deepseek", enable_thinking_kwargs={}, control_family="qwen35_xml", ), ModelSpec( slug="gemma4_12b", display_name="google/gemma-4-12B-it", hf_repo="google/gemma-4-12B-it", gguf_repo="unsloth/gemma-4-12b-it-GGUF", gguf_file="gemma-4-12b-it-Q4_K_M.gguf", gguf_arch="gemma4", num_layers=48, hidden_size=3840, intermediate_size=15360, tied_embeddings=True, rmsnorm_convention="w", mlp_activation="geglu", mid_layers=[11, 17, 23, 29], late_layers=[35, 41, 45, 46, 47], embed_key="model.language_model.embed_tokens.weight", lm_head_key=None, norm_f_key="model.language_model.norm.weight", layer_norm_tmpl="model.language_model.layers.{l}.pre_feedforward_layernorm.weight", post_ffn_norm_tmpl="model.language_model.layers.{l}.post_feedforward_layernorm.weight", gate_proj_tmpl="model.language_model.layers.{l}.mlp.gate_proj.weight", up_proj_tmpl="model.language_model.layers.{l}.mlp.up_proj.weight", down_proj_tmpl="model.language_model.layers.{l}.mlp.down_proj.weight", reasoning_format=None, enable_thinking_kwargs={"enable_thinking": True}, control_family="gemma4", ), ]