models/model_registry.py

"""KAIdol 학습 모델 레지스트리 - 모든 학습된 모델 정의"""

from typing import Dict, List, Optional

# 기본 모델 정보 (HuggingFace Hub)
BASE_MODELS = {
    "hyperclovax-32b": "naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
    "qwen2.5-72b": "Qwen/Qwen2.5-72B-Instruct",
    "qwen2.5-32b": "Qwen/Qwen2.5-32B-Instruct",
    "qwen2.5-14b": "Qwen/Qwen2.5-14B-Instruct",
    "qwen2.5-7b": "Qwen/Qwen2.5-7B-Instruct",
    "qwen3-8b": "Qwen/Qwen3-8B",
    "exaone-7.8b": "LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct",
    "solar-10.7b": "upstage/SOLAR-10.7B-Instruct-v1.0",
    "solar-pro": "upstage/solar-pro-preview-instruct",
    "varco-8b": "NCSOFT/Llama-VARCO-8B-Instruct",
    "kanana-2-30b-thinking": "kakaocorp/kanana-2-30b-a3b-thinking",
    "kanana-2-30b-instruct": "kakaocorp/kanana-2-30b-a3b-instruct",
    "llama-3.3-70b": "meta-llama/Llama-3.3-70B-Instruct",
}

# 전체 모델 레지스트리
MODEL_REGISTRY: Dict[str, Dict[str, Dict]] = {
    # ============================================================
    # DPO v5 계열 (2026-01-13)
    # ============================================================
    "dpo-v5": {
        "hyperclovax-32b-dpo-v5": {
            "path": "outputs/dpo_v5/hyperclovax-32b-dpo-v5-20260113-0012",
            "base": BASE_MODELS["hyperclovax-32b"],
            "method": "DPO",
            "size": "32B",
            "description": "HyperCLOVAX 32B DPO v5 (Primary)",
            "recommended": True,
        },
        "qwen2.5-14b-dpo-v5": {
            "path": "outputs/dpo_v5/qwen2.5-14b-dpo-v5-20260113-0045",
            "base": BASE_MODELS["qwen2.5-14b"],
            "method": "DPO",
            "size": "14B",
            "description": "Qwen2.5 14B DPO v5",
        },
        "qwen2.5-7b-dpo-v5": {
            "path": "outputs/dpo_v5/qwen2.5-7b-dpo-v5-20260113-0052",
            "base": BASE_MODELS["qwen2.5-7b"],
            "method": "DPO",
            "size": "7B",
            "description": "Qwen2.5 7B DPO v5",
        },
        "exaone-7.8b-dpo-v5": {
            "path": "outputs/dpo_v5/exaone-7.8b-dpo-v5-20260113-0052",
            "base": BASE_MODELS["exaone-7.8b"],
            "method": "DPO",
            "size": "7.8B",
            "description": "EXAONE 7.8B DPO v5",
        },
        "qwen3-8b-dpo-v5": {
            "path": "outputs/dpo_v5/qwen3-8b-dpo-v5-20260113-0052",
            "base": BASE_MODELS["qwen3-8b"],
            "method": "DPO",
            "size": "8B",
            "description": "Qwen3 8B DPO v5",
        },
        "solar-10.7b-dpo-v5": {
            "path": "outputs/dpo_v5/solar-10.7b-dpo-v5-20260113-0045",
            "base": BASE_MODELS["solar-10.7b"],
            "method": "DPO",
            "size": "10.7B",
            "description": "Solar 10.7B DPO v5",
        },
    },

    # ============================================================
    # SFT Thinking 계열 (2026-01-16)
    # ============================================================
    "sft-thinking": {
        "qwen2.5-14b-thinking": {
            "path": "outputs/qwen2.5-14b-thinking-full",
            "base": BASE_MODELS["qwen2.5-14b"],
            "method": "SFT",
            "size": "14B",
            "description": "Qwen2.5 14B SFT Thinking",
        },
        "qwen2.5-7b-thinking": {
            "path": "outputs/qwen2.5-7b-thinking-full",
            "base": BASE_MODELS["qwen2.5-7b"],
            "method": "SFT",
            "size": "7B",
            "description": "Qwen2.5 7B SFT Thinking",
        },
        "exaone-7.8b-thinking": {
            "path": "outputs/exaone-7.8b-thinking-full",
            "base": BASE_MODELS["exaone-7.8b"],
            "method": "SFT",
            "size": "7.8B",
            "description": "EXAONE 7.8B SFT Thinking",
        },
    },

    # ============================================================
    # Phase 7 Students (Kimi K2 Distillation)
    # ============================================================
    "phase7-students": {
        "kanana-30b-thinking-kimi": {
            "path": "outputs/phase7_students/kanana-2-30b-thinking-kimi-student",
            "base": BASE_MODELS["kanana-2-30b-thinking"],
            "method": "Distillation",
            "size": "30B (3B active)",
            "description": "Kanana 30B Thinking Kimi Student",
        },
        "kanana-30b-instruct-kimi": {
            "path": "outputs/phase7_students/kanana-2-30b-instruct-kimi-student",
            "base": BASE_MODELS["kanana-2-30b-instruct"],
            "method": "Distillation",
            "size": "30B (3B active)",
            "description": "Kanana 30B Instruct Kimi Student",
        },
        "qwen2.5-14b-kimi": {
            "path": "outputs/phase7_students/qwen2.5-14b-kimi-student",
            "base": BASE_MODELS["qwen2.5-14b"],
            "method": "Distillation",
            "size": "14B",
            "description": "Qwen2.5 14B Kimi Student",
        },
        "qwen2.5-7b-kimi-v3": {
            "path": "outputs/phase7_students/qwen2.5-7b-kimi-student-v3",
            "base": BASE_MODELS["qwen2.5-7b"],
            "method": "Distillation",
            "size": "7B",
            "description": "Qwen2.5 7B Kimi Student v3",
        },
        "exaone-7.8b-kimi": {
            "path": "outputs/phase7_students/exaone-7.8b-kimi-student",
            "base": BASE_MODELS["exaone-7.8b"],
            "method": "Distillation",
            "size": "7.8B",
            "description": "EXAONE 7.8B Kimi Student",
        },
    },

    # ============================================================
    # V7 Students (Latest - 2026-01-17~19)
    # ============================================================
    "v7-students": {
        "qwen2.5-72b-v7": {
            "path": "outputs/v7_students/qwen2.5-72b-v7-20260119-1113",
            "base": BASE_MODELS["qwen2.5-72b"],
            "method": "SFT",
            "size": "72B",
            "description": "Qwen2.5 72B V7 (Latest)",
        },
        "llama-3.3-70b-v7": {
            "path": "outputs/v7_students/llama-3.3-70b-v7-20260119-1114",
            "base": BASE_MODELS["llama-3.3-70b"],
            "method": "SFT",
            "size": "70B",
            "description": "Llama 3.3 70B V7 (Latest)",
        },
        "qwen2.5-32b-v7": {
            "path": "outputs/v7_students/qwen2.5-32b-v7-20260118-1135",
            "base": BASE_MODELS["qwen2.5-32b"],
            "method": "SFT",
            "size": "32B",
            "description": "Qwen2.5 32B V7",
        },
        "qwen2.5-14b-v7": {
            "path": "outputs/v7_students/qwen2.5-14b-v7-20260118-1135",
            "base": BASE_MODELS["qwen2.5-14b"],
            "method": "SFT",
            "size": "14B",
            "description": "Qwen2.5 14B V7",
        },
        "qwen2.5-7b-v7": {
            "path": "outputs/v7_students/qwen2.5-7b-v7-20260118-1135",
            "base": BASE_MODELS["qwen2.5-7b"],
            "method": "SFT",
            "size": "7B",
            "description": "Qwen2.5 7B V7",
        },
        "exaone-7.8b-v7": {
            "path": "outputs/v7_students/exaone-7.8b-v7-20260118-1135",
            "base": BASE_MODELS["exaone-7.8b"],
            "method": "SFT",
            "size": "7.8B",
            "description": "EXAONE 7.8B V7",
        },
        "qwen3-8b-v7": {
            "path": "outputs/v7_students/qwen3-8b-v7-20260118-1135",
            "base": BASE_MODELS["qwen3-8b"],
            "method": "SFT",
            "size": "8B",
            "description": "Qwen3 8B V7",
        },
        "solar-pro-v7": {
            "path": "outputs/v7_students/solar-pro-v7-20260118-1135",
            "base": BASE_MODELS["solar-pro"],
            "method": "SFT",
            "size": "22B",
            "description": "Solar Pro V7",
        },
        "varco-8b-v7": {
            "path": "outputs/v7_students/varco-8b-v7-20260118-1135",
            "base": BASE_MODELS["varco-8b"],
            "method": "SFT",
            "size": "8B",
            "description": "VARCO 8B V7",
        },
    },

    # ============================================================
    # 기타 학습 모델 (DPO, etc.)
    # ============================================================
    "others": {
        "exaone-7.8b-dpo": {
            "path": "outputs/exaone-7.8b-dpo",
            "base": BASE_MODELS["exaone-7.8b"],
            "method": "DPO",
            "size": "7.8B",
            "description": "EXAONE 7.8B DPO (Standalone)",
        },
        "qwen2.5-7b-dpo": {
            "path": "outputs/qwen2.5-7b-dpo",
            "base": BASE_MODELS["qwen2.5-7b"],
            "method": "DPO",
            "size": "7B",
            "description": "Qwen2.5 7B DPO (Standalone)",
        },
    },
}


def get_all_models() -> List[str]:
    """모든 모델 ID 목록 반환"""
    models = []
    for category, model_dict in MODEL_REGISTRY.items():
        models.extend(model_dict.keys())
    return models


def get_model_info(model_id: str) -> Optional[Dict]:
    """모델 ID로 정보 조회"""
    for category, model_dict in MODEL_REGISTRY.items():
        if model_id in model_dict:
            info = model_dict[model_id].copy()
            info["category"] = category
            info["id"] = model_id
            return info
    return None


def get_models_by_category(category: str) -> List[str]:
    """카테고리별 모델 목록"""
    return list(MODEL_REGISTRY.get(category, {}).keys())


def get_all_categories() -> List[str]:
    """모든 카테고리 목록"""
    return list(MODEL_REGISTRY.keys())


def get_models_for_dropdown() -> List[tuple]:
    """드롭다운용 (display_name, model_id) 튜플 리스트"""
    result = []
    for category, model_dict in MODEL_REGISTRY.items():
        for model_id, info in model_dict.items():
            display = f"[{info.get('size', '?')}] {info.get('description', model_id)}"
            result.append((display, model_id))
    return result


def get_small_models(max_size_gb: int = 16) -> List[str]:
    """메모리 제한에 맞는 소형 모델만 반환 (4bit 양자화 기준)"""
    # 4bit 양자화 시 대략적인 메모리: 7B~2GB, 14B~4GB, 32B~8GB, 72B~18GB
    size_map = {
        "7B": 2, "7.8B": 2, "8B": 2,
        "10.7B": 3, "14B": 4, "22B": 6,
        "30B (3B active)": 1,  # MoE
        "32B": 8, "70B": 18, "72B": 18,
    }

    result = []
    for model_id in get_all_models():
        info = get_model_info(model_id)
        if info:
            size_str = info.get("size", "72B")
            estimated_gb = size_map.get(size_str, 20)
            if estimated_gb <= max_size_gb:
                result.append(model_id)
    return result