providers.py

"""Embedding model sağlayıcıları (plug-in).

Yeni bir model eklemek = yeni bir `EmbedderProvider` alt sınıfı (gerekirse) +
`REGISTRY`'ye tek satır. Ağır kütüphaneler (torch, sentence_transformers,
fastembed) bilinçli olarak LAZY import edilir — sadece ilgili provider
örneklendiğinde yüklenir. Böylece `python -m eval --help`, ruff ve mock-tabanlı
testler model indirmeden / torch kurulmadan çalışır.

Cosine = dot product konvansiyonu: her provider çıktısını L2-normalize eder,
böylece retrieve tarafı sadece nokta çarpımı yapar.
"""

from __future__ import annotations

import os
from abc import ABC, abstractmethod
from collections.abc import Callable, Sequence

import numpy as np

# Qwen3-Embedding asimetrik kullanılır: query tarafına bir "instruction"
# eklenir, doküman tarafına eklenmez. Resmî format: "Instruct: {görev}\nQuery: ".
# (Türkçe corpus'ta İngilizce instruction sorunsuz çalışır — model çok dillidir.)
QWEN3_QUERY_INSTRUCTION = (
    "Instruct: Given a search query, retrieve relevant passages that answer "
    "the query\nQuery: "
)


def l2_normalize(vecs: np.ndarray) -> np.ndarray:
    """Satır bazında L2-normalize (sıfır vektörlere karşı korumalı)."""
    vecs = np.asarray(vecs, dtype=np.float32)
    norms = np.linalg.norm(vecs, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    return vecs / norms


def _auto_device() -> str:
    import torch

    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


class EmbedderProvider(ABC):
    """Tüm modellerin uyduğu arayüz.

    Alt sınıflar `name` (cache/tablo etiketi) ve `dim` (vektör boyutu) alanlarını
    doldurur; corpus ve query encode'unu uygular. Çıktılar L2-normalize'li
    np.float32 matrisleridir (shape: [n, dim]).
    """

    name: str
    dim: int

    @abstractmethod
    def encode_corpus(self, texts: Sequence[str]) -> np.ndarray: ...

    @abstractmethod
    def encode_queries(self, texts: Sequence[str]) -> np.ndarray: ...


class SentenceTransformerProvider(EmbedderProvider):
    """sentence-transformers ile yüklenen modeller (bge-m3, Qwen3-Embedding).

    `query_prompt` verilirse SADECE query encode'unda her metnin başına eklenir
    (Qwen3 instruction'ı için). Doküman tarafı prompt'suz kalır. bge-m3 için
    `query_prompt=None` -> simetrik kullanım.
    """

    def __init__(
        self,
        model_id: str,
        name: str | None = None,
        query_prompt: str | None = None,
        batch_size: int = 16,
    ) -> None:
        from sentence_transformers import SentenceTransformer

        self.model_id = model_id
        self.name = name or model_id.split("/")[-1]
        self._query_prompt = query_prompt
        self.batch_size = batch_size
        self.model = SentenceTransformer(model_id, device=_auto_device())
        # ST 5.x'te get_sentence_embedding_dimension -> get_embedding_dimension.
        get_dim = getattr(self.model, "get_embedding_dimension", None) or (
            self.model.get_sentence_embedding_dimension
        )
        self.dim = int(get_dim())

    def _encode(self, texts: Sequence[str], prompt: str | None) -> np.ndarray:
        embs = self.model.encode(
            list(texts),
            batch_size=self.batch_size,
            prompt=prompt,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
        )
        return np.asarray(embs, dtype=np.float32)

    def encode_corpus(self, texts: Sequence[str]) -> np.ndarray:
        return self._encode(texts, prompt=None)

    def encode_queries(self, texts: Sequence[str]) -> np.ndarray:
        return self._encode(texts, prompt=self._query_prompt)


class FastEmbedProvider(EmbedderProvider):
    """fastembed (ONNX) ile yüklenen modeller — prod baseline (BGE-Base 768d).

    fastembed BGE çıktısını normalize eder; yine de tutarlılık için L2-normalize
    uygularız. Bu provider'ın amacı: mevcut prod modelinin Türkçe corpus'taki
    skorunu yeni adaylarla aynı tabloda görmek (migrasyon gerekçesi).
    """

    def __init__(self, model_id: str, name: str | None = None, dim: int = 768) -> None:
        from fastembed import TextEmbedding

        self.model_id = model_id
        self.name = name or model_id.split("/")[-1]
        self.dim = dim
        self.model = TextEmbedding(model_name=model_id)

    def _encode(self, texts: Sequence[str]) -> np.ndarray:
        embs = np.array(list(self.model.embed(list(texts))), dtype=np.float32)
        self.dim = int(embs.shape[1])
        return l2_normalize(embs)

    def encode_corpus(self, texts: Sequence[str]) -> np.ndarray:
        return self._encode(texts)

    def encode_queries(self, texts: Sequence[str]) -> np.ndarray:
        return self._encode(texts)


class OllamaProvider(EmbedderProvider):
    """Yerel Ollama servisinden embedding alır (HTTP /api/embed, stdlib urllib).

    Avantaj: model zaten `ollama pull` ile çekilmişse yeniden indirme yok ve
    kurulum hafiftir (torch/transformers gerekmez). bge-m3 Ollama'da F16'dır;
    embedding sıralaması açısından fp32'ye göre farkı ihmal edilebilir.

    Host: OLLAMA_HOST env değişkeni ya da http://localhost:11434.
    """

    def __init__(
        self,
        model_id: str,
        name: str | None = None,
        dim: int = 1024,
        host: str | None = None,
        batch_size: int = 32,
    ) -> None:
        self.model_id = model_id
        self.name = name or f"ollama-{model_id}"
        self.dim = dim
        self.batch_size = batch_size
        host = host or os.environ.get("OLLAMA_HOST") or "http://localhost:11434"
        if not host.startswith("http"):
            host = f"http://{host}"
        self.host = host.rstrip("/")

    def _embed(self, texts: Sequence[str]) -> np.ndarray:
        import json
        import urllib.request

        texts = list(texts)
        out: list[list[float]] = []
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i : i + self.batch_size]
            payload = json.dumps({"model": self.model_id, "input": batch}).encode("utf-8")
            req = urllib.request.Request(
                f"{self.host}/api/embed",
                data=payload,
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=300) as resp:
                data = json.loads(resp.read())
            out.extend(data["embeddings"])
        arr = np.asarray(out, dtype=np.float32)
        self.dim = int(arr.shape[1])
        return l2_normalize(arr)

    def encode_corpus(self, texts: Sequence[str]) -> np.ndarray:
        return self._embed(texts)

    def encode_queries(self, texts: Sequence[str]) -> np.ndarray:
        return self._embed(texts)


class MockEmbedder(EmbedderProvider):
    """İndirme gerektirmeyen, deterministik sahte embedder.

    Kelime token'larını sabit bir hash ile boyutlara dağıtan bag-of-words
    vektörü üretir. Gerçek anlamı yoktur; SADECE harness plumbing'ini (embed ->
    retrieve -> score akışı) hızlı/çevrimdışı doğrulamak içindir. CI smoke gerçek
    modelleri kullanır; bu mock yerel hızlı kontrol için elde durur.
    """

    def __init__(self, name: str = "mock", dim: int = 64) -> None:
        self.name = name
        self.dim = dim

    def _encode(self, texts: Sequence[str]) -> np.ndarray:
        import hashlib

        vecs = np.zeros((len(list(texts)), self.dim), dtype=np.float32)
        for i, t in enumerate(texts):
            for tok in t.split():
                h = int.from_bytes(hashlib.md5(tok.encode("utf-8")).digest()[:4], "little")
                vecs[i, h % self.dim] += 1.0
        return l2_normalize(vecs)

    def encode_corpus(self, texts: Sequence[str]) -> np.ndarray:
        return self._encode(texts)

    def encode_queries(self, texts: Sequence[str]) -> np.ndarray:
        return self._encode(texts)


# --- Plug-in kayıt defteri --------------------------------------------------
# Yeni model = buraya tek satır. Değerler "factory" (lambda); model ancak
# çağrıldığında yüklenir, bu yüzden REGISTRY'yi import etmek ağır değildir.
#
# NOT: baseline model id'si prod ile birebir aynı olmalı. enterprise-ai-assistant
# tarafındaki fastembed model adını teyit edip gerekirse burayı güncelle.
REGISTRY: dict[str, Callable[[], EmbedderProvider]] = {
    "bge-m3": lambda: SentenceTransformerProvider("BAAI/bge-m3", name="bge-m3"),
    "qwen3-0.6b": lambda: SentenceTransformerProvider(
        "Qwen/Qwen3-Embedding-0.6B",
        name="qwen3-0.6b",
        query_prompt=QWEN3_QUERY_INSTRUCTION,
    ),
    "baseline-bge-base": lambda: FastEmbedProvider(
        "BAAI/bge-base-en-v1.5", name="baseline-bge-base", dim=768
    ),
    # Yerel Ollama'dan bge-m3 (F16). Kurulu/çekili ise indirme gerektirmez.
    "bge-m3-ollama": lambda: OllamaProvider("bge-m3", name="bge-m3-ollama", dim=1024),
    "mock": lambda: MockEmbedder(),
}

# Tablo/karşılaştırmada baseline olarak kabul edilen model (delta'lar buna göre).
BASELINE_MODEL = "baseline-bge-base"

# `--models` verilmediğinde koşulan gerçek modeller (mock hariç).
DEFAULT_MODELS = ["bge-m3", "qwen3-0.6b", "baseline-bge-base"]


def get_provider(name: str) -> EmbedderProvider:
    """REGISTRY'den bir provider örneği üretir (modeli burada yükler)."""
    if name not in REGISTRY:
        raise KeyError(
            f"Bilinmeyen model '{name}'. Kayıtlı modeller: {', '.join(REGISTRY)}"
        )
    return REGISTRY[name]()