live-app-translator/src/livesub/models/tiers.py

"""모델 품질 티어 정의.

사용자에게는 실제 모델 이름 대신 5단계 티어(속도우선 → 품질우선)로 노출한다.
각 티어는 ASR(음성인식) 1개 + MT(번역) 1개의 조합이다.

선정 근거는 docs/MODELS.md 참고.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum

from ..constants import DEFAULT_TIER


class MTBackend(str, Enum):
    """번역 모델 실행 백엔드."""

    CTRANSLATE2 = "ctranslate2"  # NLLB 계열 seq2seq
    TRANSFORMERS = "transformers"  # LLM 계열 (Seed-X / Qwen3)


class PromptStyle(str, Enum):
    """LLM 백엔드의 프롬프트 형식. 모델 계열마다 요구사항이 다르다."""

    #: 해당 없음 (seq2seq)
    NONE = "none"
    #: Seed-X 전용. chat template 이 없는 번역 전용 completion 모델이라
    #: 모델 카드에 적힌 고정 형식 + 끝의 `<언어코드>` 태그를 그대로 지켜야 한다.
    SEEDX = "seedx"
    #: Qwen3 등 일반 instruct 모델. chat template + 자유 지시문 사용 가능.
    INSTRUCT = "instruct"


@dataclass(frozen=True)
class AsrSpec:
    repo: str
    compute_type: str
    beam_size: int = 1
    vram_mb: int = 0


@dataclass(frozen=True)
class MTSpec:
    repo: str
    backend: MTBackend
    compute_type: str = "float16"
    vram_mb: int = 0
    prompt_style: PromptStyle = PromptStyle.NONE
    #: LoRA 추가학습 난이도 1(쉬움) ~ 3(까다로움)
    finetune_ease: int = 2

    @property
    def supports_prompt_glossary(self) -> bool:
        """프롬프트로 용어집을 지시할 수 있는가.

        지시문을 이해하는 instruct 모델만 가능하다. Seed-X 는 번역만 하도록
        학습된 completion 모델이라 "이 용어는 이렇게 옮겨라" 같은 문장을
        넣으면 학습 분포를 벗어나 오히려 번역이 망가진다. seq2seq 도 마찬가지다.
        이 둘은 플레이스홀더 치환 방식으로 용어를 보호한다.
        """
        return self.prompt_style is PromptStyle.INSTRUCT


@dataclass(frozen=True)
class Tier:
    key: str
    order: int
    name: str
    tagline: str
    asr: AsrSpec
    mt: MTSpec
    approx_latency_s: float
    min_vram_gb: float
    notes: str = ""
    recommended: bool = False
    best_after_finetune: bool = False

    @property
    def total_vram_mb(self) -> int:
        return self.asr.vram_mb + self.mt.vram_mb

    @property
    def download_gb(self) -> float:
        return round(self.total_vram_mb / 1024 * 1.25, 1)


TIERS: dict[str, Tier] = {}


def _register(tier: Tier) -> Tier:
    TIERS[tier.key] = tier
    return tier


LIGHTNING = _register(
    Tier(
        key="lightning",
        order=1,
        name="번개",
        tagline="속도 최우선 · 저사양 GPU",
        asr=AsrSpec("Systran/faster-whisper-small", "int8_float16", 1, 700),
        mt=MTSpec(
            "entai2965/nllb-200-distilled-600M-ctranslate2",
            MTBackend.CTRANSLATE2,
            "int8_float16",
            800,
            finetune_ease=1,
        ),
        approx_latency_s=0.6,
        min_vram_gb=2,
        notes="짧은 대사·게임 음성에 충분. 긴 문장이나 전문용어는 품질이 떨어질 수 있습니다.",
    )
)

SWIFT = _register(
    Tier(
        key="swift",
        order=2,
        name="신속",
        tagline="빠르면서 인식률은 최상급",
        asr=AsrSpec("Systran/faster-whisper-large-v3-turbo", "int8_float16", 1, 1600),
        mt=MTSpec(
            "entai2965/nllb-200-distilled-600M-ctranslate2",
            MTBackend.CTRANSLATE2,
            "float16",
            1400,
            finetune_ease=1,
        ),
        approx_latency_s=0.9,
        min_vram_gb=4,
        notes="음성인식을 최상급으로 올리고 번역만 경량으로 유지한 구성. 가성비가 가장 좋습니다.",
    )
)

BALANCE = _register(
    Tier(
        key="balance",
        order=3,
        name="균형",
        tagline="기본값 · 속도와 품질의 중간",
        asr=AsrSpec("Systran/faster-whisper-large-v3-turbo", "float16", 2, 2400),
        mt=MTSpec(
            "entai2965/nllb-200-distilled-1.3B-ctranslate2",
            MTBackend.CTRANSLATE2,
            "float16",
            2900,
            finetune_ease=2,
        ),
        approx_latency_s=1.2,
        min_vram_gb=6,
        notes="8GB VRAM 환경에서 가장 무난한 선택.",
    )
)

PRECISION = _register(
    Tier(
        key="precision",
        order=4,
        name="정밀",
        tagline="추천 · 현 시점 최고 번역 품질",
        asr=AsrSpec("Systran/faster-whisper-large-v3-turbo", "float16", 2, 2400),
        mt=MTSpec(
            "ByteDance-Seed/Seed-X-PPO-7B-AWQ-Int4",
            MTBackend.TRANSFORMERS,
            "int4",
            5600,
            prompt_style=PromptStyle.SEEDX,
            finetune_ease=3,
        ),
        approx_latency_s=2.0,
        min_vram_gb=10,
        notes="번역 전용으로 학습된 7B 모델. 한/영/일/중 구어체와 문맥 처리가 확연히 좋습니다.",
        recommended=True,
    )
)

ULTIMATE = _register(
    Tier(
        key="ultimate",
        order=5,
        name="극한",
        tagline="품질 최우선 · 추가학습(LoRA)에 최적",
        asr=AsrSpec("Systran/faster-whisper-large-v3", "float16", 5, 3100),
        mt=MTSpec(
            "Qwen/Qwen3-8B",
            MTBackend.TRANSFORMERS,
            "float16",
            11000,
            prompt_style=PromptStyle.INSTRUCT,
            finetune_ease=1,
        ),
        approx_latency_s=3.0,
        min_vram_gb=16,
        notes="게임·방송 용어를 LoRA로 추가학습해 붙일 때 가장 잘 먹히는 구성입니다.",
        best_after_finetune=True,
    )
)

assert DEFAULT_TIER in TIERS, f"constants.DEFAULT_TIER({DEFAULT_TIER}) 가 티어 목록에 없습니다"


def ordered_tiers() -> list[Tier]:
    return sorted(TIERS.values(), key=lambda t: t.order)


def get_tier(key: str) -> Tier:
    return TIERS.get(key, TIERS[DEFAULT_TIER])