model_families.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Family-based model resolution for LLMs beyond the catalog."""
 
from __future__ import annotations
 
import json
import re
from pathlib import Path
from typing import Any
 
_RUNTIME_LLM: dict[str, dict[str, Any]] = {}
FAMILY_MODEL_DIR = Path.home() / ".aquin" / "family_models"
 
# HF AutoConfig.model_type → family id (fine-tunes under any org)
_MODEL_TYPE_TO_FAMILY: dict[str, str] = {
    "gpt2": "gpt2",
    "gpt_neox": "pythia",
    "llama": "llama",
    "mistral": "llama",  # same LoRA targets / hook shape as Llama
    "qwen": "qwen",
    "qwen2": "qwen",
    "qwen3": "qwen",
    "qwen2_moe": "qwen",
}
 
# Architecture class name substrings → family id (backup when model_type is missing)
_ARCH_HINTS: tuple[tuple[str, str], ...] = (
    ("Qwen3", "qwen"),
    ("Qwen2", "qwen"),
    ("Qwen", "qwen"),
    ("Llama", "llama"),
    ("Mistral", "llama"),
    ("GPTNeoX", "pythia"),
    ("GPT2", "gpt2"),
    ("LFM", "lfm"),
)
 
LLM_FAMILIES: list[dict[str, Any]] = [
    {
        "id": "gpt2",
        "label": "GPT-2",
        "hf_patterns": (re.compile(r"^(?:openai/)?gpt2(?:-|$)", re.I),),
        "lora_target_modules": ["c_attn", "c_fc"],
        "hook": "blocks.{layer}.hook_resid_post",
        "examples": ["gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"],
        "trust_remote_code": False,
    },
    {
        "id": "pythia",
        "label": "Pythia",
        "hf_patterns": (re.compile(r"^EleutherAI/pythia", re.I),),
        "slug_prefix": "EleutherAI/",
        "lora_target_modules": ["query_key_value", "dense"],
        "hook": "blocks.{layer}.hook_resid_post",
        "examples": [
            "EleutherAI/pythia-160m-deduped",
            "EleutherAI/pythia-410m-deduped",
            "EleutherAI/pythia-1.4b-deduped",
            "EleutherAI/pythia-2.8b",
        ],
        "trust_remote_code": False,
    },
    {
        "id": "llama",
        "label": "Llama",
        "hf_patterns": (
            re.compile(r"^meta-llama/Llama", re.I),
            re.compile(r"^meta-llama/Meta-Llama", re.I),
        ),
        "lora_target_modules": ["q_proj", "v_proj"],
        "hook": "blocks.{layer}.hook_resid_post",
        "examples": [
            "meta-llama/Llama-3.2-1B-Instruct",
            "meta-llama/Llama-3.2-3B-Instruct",
            "meta-llama/Llama-3.1-8B-Instruct",
        ],
        "trust_remote_code": False,
    },
    {
        "id": "lfm",
        "label": "LFM (Liquid AI)",
        "hf_patterns": (re.compile(r"^LiquidAI/LFM", re.I),),
        "lora_target_modules": ["q_proj", "v_proj"],
        "hook": "blocks.{layer}.hook_resid_post",
        "examples": [
            "LiquidAI/LFM2.5-1.2B-Instruct",
            "LiquidAI/LFM2.5-1.2B-Thinking",
            "LiquidAI/LFM2.5-230M",
        ],
        "trust_remote_code": False,
        "hf_only": True,
    },
    {
        "id": "qwen",
        "label": "Qwen",
        # Official Qwen/* plus community fine-tunes with "qwen" in the repo name
        "hf_patterns": (
            re.compile(r"^Qwen/Qwen", re.I),
            re.compile(r"(?i)(?:^|/)[^/\s]*qwen"),
        ),
        "lora_target_modules": ["q_proj", "v_proj"],
        "hook": "blocks.{layer}.hook_resid_post",
        "examples": [
            "Qwen/Qwen3-8B",
            "Qwen/Qwen2.5-7B-Instruct",
            "Qwen/Qwen2.5-0.5B-Instruct",
        ],
        "trust_remote_code": False,
        "hf_only": True,
    },
]
 
 
 
def _slug_from_hf(hf_name: str) -> str:
    """Stable Aquin slug from a HuggingFace repo id."""
    part = hf_name.split("/")[-1].lower()
    part = re.sub(r"[^a-z0-9._-]+", "-", part).strip("-")
    return part or "model"
 
 
def _match_family(hf_name: str, families: list[dict[str, Any]]) -> dict[str, Any] | None:
    for fam in families:
        for pat in fam.get("hf_patterns", ()):
            if pat.search(hf_name):
                return fam
    return None
 
 
def _family_by_id(family_id: str) -> dict[str, Any] | None:
    for fam in LLM_FAMILIES:
        if fam["id"] == family_id:
            return fam
    return None
 
 
def _match_family_from_hf_config(hf_name: str) -> dict[str, Any] | None:
    """Infer LLM family from HuggingFace config when the repo id is not an official org prefix.
 
    Covers fine-tunes (e.g. Hahmdong/*-qwen3-*) that keep the base architecture in config.json.
    """
    try:
        from transformers import AutoConfig
 
        cfg = AutoConfig.from_pretrained(hf_name, trust_remote_code=True)
    except Exception:
        return None
 
    model_type = str(getattr(cfg, "model_type", "") or "").lower()
    fam_id = _MODEL_TYPE_TO_FAMILY.get(model_type)
    if fam_id:
        return _family_by_id(fam_id)
 
    arches = getattr(cfg, "architectures", None) or []
    for arch in arches:
        name = str(arch)
        for hint, fid in _ARCH_HINTS:
            if hint.lower() in name.lower():
                return _family_by_id(fid)
    return None
 
 
def normalize_llm_hf_id(model_id: str) -> str | None:
    raw = (model_id or "").strip()
    if not raw:
        return None
    if "/" in raw:
        return raw
    if re.match(r"^gpt2(?:-|$)", raw, re.I):
        return raw
    if raw.lower().startswith("pythia"):
        return f"EleutherAI/{raw}"
    return None
 
 
 
def _infer_llm_dims(hf_name: str, *, trust_remote_code: bool) -> tuple[int, int]:
    from transformers import AutoConfig
 
    cfg = AutoConfig.from_pretrained(hf_name, trust_remote_code=trust_remote_code)
    n_layers = int(
        getattr(cfg, "num_hidden_layers", None)
        or getattr(cfg, "n_layer", None)
        or getattr(cfg, "n_layers", 0)
    )
    d_model = int(
        getattr(cfg, "hidden_size", None)
        or getattr(cfg, "n_embd", None)
        or getattr(cfg, "d_model", 0)
    )
    if n_layers < 1 or d_model < 1:
        raise ValueError(f"Could not infer layer count / width for {hf_name}")
    return d_model, n_layers
 
 
 
def build_llm_family_config(hf_name: str, family: dict[str, Any]) -> dict[str, Any]:
    from transformers import AutoConfig
 
    trust = bool(family.get("trust_remote_code", False))
    d_model, n_layers = _infer_llm_dims(hf_name, trust_remote_code=trust)
    hcfg = AutoConfig.from_pretrained(hf_name, trust_remote_code=trust)
    n_heads = int(
        getattr(hcfg, "num_attention_heads", None)
        or getattr(hcfg, "n_head", None)
        or 32
    )
    sae_layer = max(0, n_layers // 2)
    slug = _slug_from_hf(hf_name)
    cfg: dict[str, Any] = {
        "hf_name": hf_name,
        "d_model": d_model,
        "n_layers": n_layers,
        "n_heads": n_heads,
        "sae_layer": sae_layer,
        "sae_source": "user",
        "sae_layers": {},
        "family": family["id"],
        "hook_template": family.get("hook", "blocks.{layer}.hook_resid_post"),
        "lora_target_modules": list(family.get("lora_target_modules", ["q_proj", "v_proj"])),
        "trust_remote_code": trust,
        "catalog": False,
        "_slug": slug,
    }
    if family.get("hf_only"):
        cfg["hf_only"] = True
    if family.get("hf_first"):
        cfg["hf_first"] = True
    return cfg
 
 
 
def _persist_family_config(slug: str, cfg: dict[str, Any]) -> None:
    FAMILY_MODEL_DIR.mkdir(parents=True, exist_ok=True)
    payload = {k: v for k, v in cfg.items() if not str(k).startswith("_")}
    path = FAMILY_MODEL_DIR / f"{slug}.json"
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
 
 
def _load_family_config_from_disk(slug: str) -> dict[str, Any] | None:
    path = FAMILY_MODEL_DIR / f"{slug}.json"
    if not path.is_file():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
    return data if isinstance(data, dict) else None
 
 
def get_runtime_llm_config(slug: str) -> dict[str, Any] | None:
    if slug in _RUNTIME_LLM:
        return _RUNTIME_LLM[slug]
    cfg = _load_family_config_from_disk(slug)
    if cfg:
        _RUNTIME_LLM[slug] = cfg
        return cfg
    return None
 
 
def resolve_llm_family(model_id: str) -> tuple[str, dict[str, Any]]:
    """Resolve catalog-external LLM slug or HF id via family rules. Raises ValueError if unsupported."""
    raw = (model_id or "").strip()
    if not raw:
        raise ValueError("model_id required")
 
    cached = get_runtime_llm_config(raw)
    if cached:
        return raw, cached
 
    hf_name = normalize_llm_hf_id(raw) or raw
 
    family = _match_family(hf_name, LLM_FAMILIES)
    if not family and "/" in hf_name:
        family = _match_family_from_hf_config(hf_name)
    if not family and "/" not in hf_name:
        raise ValueError(
            f"Unknown LLM '{model_id}'. Pass a HuggingFace repo id "
            f"(e.g. gpt2-medium, EleutherAI/pythia-410m-deduped, meta-llama/Llama-3.2-3B-Instruct)."
        )
    if not family:
        raise ValueError(
            f"LLM '{hf_name}' is not in a supported family yet. "
            f"Supported: GPT-2, Pythia (EleutherAI/pythia-*), Llama (meta-llama/Llama*), "
            f"LFM (LiquidAI/LFM*), Qwen (Qwen/Qwen* and fine-tunes with qwen in the name / "
            f"qwen2|qwen3 architecture)."
        )
 
    slug = _slug_from_hf(hf_name)
    cached = get_runtime_llm_config(slug)
    if cached:
        return slug, cached
 
    cfg = build_llm_family_config(hf_name, family)
    _RUNTIME_LLM[slug] = cfg
    _persist_family_config(slug, cfg)
    return slug, cfg
 
 
 
 
def is_llm_family_candidate(model_id: str) -> bool:
    raw = (model_id or "").strip()
    if get_runtime_llm_config(raw):
        return True
    hf_name = normalize_llm_hf_id(raw) or raw
    if "/" not in hf_name and not re.match(r"^gpt2", raw, re.I):
        return False
    return _match_family(hf_name, LLM_FAMILIES) is not None
 
 
 
def family_help_lines() -> list[str]:
    lines = ["Model families (pass HuggingFace repo id to aquin load --model):", ""]
    lines.append("LLMs:")
    for fam in LLM_FAMILIES:
        ex = ", ".join(fam.get("examples", [])[:3])
        lines.append(f"  {fam['label']}: {ex}")
    lines.append("")
    lines.append("Load SAEs with: aquin load sae <model-l{n}> or --path <file.pt>")
    lines.append("Family models use inspection tools; self-train SAE: aquin sae train (see /docs/sae-training).")
    return lines