benchmarking.py

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Sequence, Tuple


@dataclass(frozen=True)
class BenchmarkModelSpec:
    key: str
    label: str
    model_id: str
    kind: str = "sentence_transformer"
    input_prefix: str = ""
    default: bool = False
    note: str = ""


BENCHMARK_MODEL_SPECS: Tuple[BenchmarkModelSpec, ...] = (
    BenchmarkModelSpec(
        key="current-adapter",
        label="현재 모델 (merchant-region adapter)",
        model_id="hoin1218/bge-m3-merchant-region-pair-adapter",
        kind="current_adapter",
        default=True,
        note="BGE-M3 base + merchant/region projection adapter",
    ),
    BenchmarkModelSpec(
        key="bge-m3",
        label="BAAI/bge-m3",
        model_id="BAAI/bge-m3",
        note="Heavy multilingual baseline",
    ),
    BenchmarkModelSpec(
        key="multilingual-e5-small",
        label="intfloat/multilingual-e5-small",
        model_id="intfloat/multilingual-e5-small",
        input_prefix="query: ",
        default=True,
        note="Fast multilingual E5 baseline",
    ),
    BenchmarkModelSpec(
        key="multilingual-minilm",
        label="paraphrase-multilingual-MiniLM-L12-v2",
        model_id="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        default=True,
        note="Compact multilingual sentence-transformers baseline",
    ),
    BenchmarkModelSpec(
        key="ko-sroberta",
        label="jhgan/ko-sroberta-multitask",
        model_id="jhgan/ko-sroberta-multitask",
        note="Korean sentence embedding baseline",
    ),
    BenchmarkModelSpec(
        key="kosimcse-roberta-multitask",
        label="BM-K/KoSimCSE-roberta-multitask",
        model_id="BM-K/KoSimCSE-roberta-multitask",
        kind="transformers_cls",
        note="Korean KoSimCSE multitask RoBERTa baseline",
    ),
    BenchmarkModelSpec(
        key="kosimcse-roberta",
        label="BM-K/KoSimCSE-roberta",
        model_id="BM-K/KoSimCSE-roberta",
        kind="transformers_cls",
        note="Korean KoSimCSE RoBERTa baseline",
    ),
    BenchmarkModelSpec(
        key="kr-sbert",
        label="KR-SBERT V40K klueNLI augSTS",
        model_id="snunlp/KR-SBERT-V40K-klueNLI-augSTS",
        note="Korean SBERT sentence-similarity baseline",
    ),
    BenchmarkModelSpec(
        key="kure-v1",
        label="nlpai-lab/KURE-v1",
        model_id="nlpai-lab/KURE-v1",
        note="Korean universal representation embedding baseline",
    ),
)

DEFAULT_BENCHMARK_MODEL_IDS: Tuple[str, ...] = tuple(
    spec.key for spec in BENCHMARK_MODEL_SPECS if spec.default
)

BENCHMARK_TABLE_HEADERS: Tuple[str, ...] = (
    "Model",
    "Similarity",
    "Judgement",
    "Elapsed ms",
    "Status",
)


def get_benchmark_choices() -> List[Tuple[str, str]]:
    return [(spec.label, spec.key) for spec in BENCHMARK_MODEL_SPECS]


def get_benchmark_specs(
    selected_model_ids: Optional[Sequence[str]],
) -> List[BenchmarkModelSpec]:
    selected = (
        set(DEFAULT_BENCHMARK_MODEL_IDS)
        if selected_model_ids is None
        else set(selected_model_ids)
    )
    return [spec for spec in BENCHMARK_MODEL_SPECS if spec.key in selected]


def prepare_benchmark_text(spec: BenchmarkModelSpec, text: str) -> str:
    normalized = (text or "").strip()
    if spec.input_prefix and not normalized.startswith(spec.input_prefix):
        return f"{spec.input_prefix}{normalized}"
    return normalized


def build_benchmark_row(
    spec: BenchmarkModelSpec,
    score: Optional[float],
    judgement: str,
    elapsed_ms: Optional[float],
    status: str,
) -> List[object]:
    score_value = None if score is None else round(score, 4)
    elapsed_value = None if elapsed_ms is None else round(elapsed_ms, 1)
    return [spec.label, score_value, judgement, elapsed_value, status]