hackathon_advisor/llama_embedding.py

from __future__ import annotations

from collections.abc import Sequence
import atexit
import json
import os
from pathlib import Path
import platform
import subprocess
import sys
import threading
from typing import Any

from hackathon_advisor.config import bool_env, int_env, optional_int_env, tri_state_env
from hackathon_advisor.data import (
    DEFAULT_EMBEDDING_MODEL_FILE,
    DEFAULT_EMBEDDING_MODEL_REPO,
)

DEFAULT_N_CTX = 2048


class LlamaCppEmbedder:
    def __init__(
        self,
        *,
        model_repo: str = DEFAULT_EMBEDDING_MODEL_REPO,
        model_file: str = DEFAULT_EMBEDDING_MODEL_FILE,
        model_path: str = "",
        n_ctx: int = DEFAULT_N_CTX,
        n_batch: int | None = None,
        n_threads: int | None = None,
        n_gpu_layers: int = 0,
        verbose: bool = False,
    ) -> None:
        self.model_repo = model_repo.strip() or DEFAULT_EMBEDDING_MODEL_REPO
        self.model_file = model_file.strip() or DEFAULT_EMBEDDING_MODEL_FILE
        self.model_path = model_path.strip()
        self.n_ctx = n_ctx
        self.n_batch = n_batch or n_ctx
        self.n_threads = n_threads
        self.n_gpu_layers = n_gpu_layers
        self.verbose = verbose
        self._model = None

    def __call__(self, text: str) -> Sequence[float]:
        return self.embed(text)

    def embed(self, text: str) -> Sequence[float]:
        model = self._ensure_model()
        return model.embed(text, normalize=True)

    def _ensure_model(self):
        if self._model is not None:
            return self._model
        from huggingface_hub import hf_hub_download
        from llama_cpp import LLAMA_POOLING_TYPE_MEAN, Llama

        model_path = self.model_path
        if not model_path:
            model_path = hf_hub_download(
                repo_id=self.model_repo,
                filename=self.model_file,
                repo_type="model",
            )
        if not Path(model_path).is_file():
            raise RuntimeError(f"llama.cpp embedding model was not found: {model_path}")
        self._model = Llama(
            model_path=model_path,
            embedding=True,
            pooling_type=LLAMA_POOLING_TYPE_MEAN,
            n_ctx=self.n_ctx,
            n_batch=self.n_batch,
            n_ubatch=self.n_batch,
            n_threads=self.n_threads,
            n_gpu_layers=self.n_gpu_layers,
            verbose=self.verbose,
        )
        return self._model


class SubprocessLlamaCppEmbedder:
    def __init__(
        self,
        *,
        model_repo: str = DEFAULT_EMBEDDING_MODEL_REPO,
        model_file: str = DEFAULT_EMBEDDING_MODEL_FILE,
        model_path: str = "",
        n_ctx: int = DEFAULT_N_CTX,
        n_batch: int | None = None,
        n_threads: int | None = None,
        n_gpu_layers: int = 0,
        verbose: bool = False,
    ) -> None:
        self.model_repo = model_repo.strip() or DEFAULT_EMBEDDING_MODEL_REPO
        self.model_file = model_file.strip() or DEFAULT_EMBEDDING_MODEL_FILE
        self.model_path = model_path.strip()
        self.n_ctx = n_ctx
        self.n_batch = n_batch or n_ctx
        self.n_threads = n_threads
        self.n_gpu_layers = n_gpu_layers
        self.verbose = verbose
        self._process: subprocess.Popen[str] | None = None
        self._request_id = 0
        self._lock = threading.Lock()
        atexit.register(self.close)

    def __call__(self, text: str) -> Sequence[float]:
        return self.embed(text)

    def embed(self, text: str) -> Sequence[float]:
        with self._lock:
            process = self._ensure_process()
            self._request_id += 1
            request_id = self._request_id
            request = json.dumps({"id": request_id, "text": text}, ensure_ascii=False)
            try:
                assert process.stdin is not None
                assert process.stdout is not None
                process.stdin.write(f"{request}\n")
                process.stdin.flush()
                line = process.stdout.readline()
            except (BrokenPipeError, OSError) as error:
                self.close()
                raise RuntimeError("llama.cpp embedding worker stopped before returning a vector.") from error
            if not line:
                returncode = process.poll()
                self.close()
                detail = f" with exit code {returncode}" if returncode is not None else ""
                raise RuntimeError(f"llama.cpp embedding worker exited{detail}.")
            try:
                response = json.loads(line)
            except json.JSONDecodeError as error:
                raise RuntimeError("llama.cpp embedding worker returned invalid JSON.") from error
            if response.get("id") != request_id:
                raise RuntimeError("llama.cpp embedding worker returned an out-of-order response.")
            if response.get("error"):
                raise RuntimeError(str(response["error"]))
            vector = response.get("vector")
            if not isinstance(vector, list):
                raise RuntimeError("llama.cpp embedding worker did not return a vector.")
            return vector

    def close(self) -> None:
        process = self._process
        self._process = None
        if process is None:
            return
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=2)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=2)

    def _ensure_process(self) -> subprocess.Popen[str]:
        if self._process is not None and self._process.poll() is None:
            return self._process
        self._process = subprocess.Popen(
            [sys.executable, "-u", "-m", "hackathon_advisor.llama_embedding", "--worker"],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=None if self.verbose else subprocess.DEVNULL,
            text=True,
            cwd=Path(__file__).resolve().parents[1],
        )
        config = json.dumps(
            {
                "model_repo": self.model_repo,
                "model_file": self.model_file,
                "model_path": self.model_path,
                "n_ctx": self.n_ctx,
                "n_batch": self.n_batch,
                "n_threads": self.n_threads,
                "n_gpu_layers": self.n_gpu_layers,
                "verbose": self.verbose,
            },
            ensure_ascii=False,
        )
        assert self._process.stdin is not None
        self._process.stdin.write(f"{config}\n")
        self._process.stdin.flush()
        return self._process


def create_llama_cpp_embedder(metadata: dict[str, Any]) -> LlamaCppEmbedder | SubprocessLlamaCppEmbedder:
    embedder_cls = SubprocessLlamaCppEmbedder if _use_subprocess_embedder() else LlamaCppEmbedder
    return embedder_cls(
        model_repo=os.environ.get(
            "ADVISOR_EMBEDDING_MODEL_REPO",
            str(metadata.get("model_repo") or DEFAULT_EMBEDDING_MODEL_REPO),
        ),
        model_file=os.environ.get(
            "ADVISOR_EMBEDDING_MODEL_FILE",
            str(metadata.get("model_file") or DEFAULT_EMBEDDING_MODEL_FILE),
        ),
        model_path=os.environ.get("ADVISOR_EMBEDDING_MODEL_PATH", ""),
        n_ctx=int_env("ADVISOR_EMBEDDING_N_CTX", DEFAULT_N_CTX, minimum=0),
        n_batch=optional_int_env("ADVISOR_EMBEDDING_BATCH"),
        n_threads=optional_int_env("ADVISOR_EMBEDDING_THREADS"),
        n_gpu_layers=int_env("ADVISOR_EMBEDDING_GPU_LAYERS", 0, minimum=0),
        verbose=bool_env("ADVISOR_EMBEDDING_VERBOSE"),
    )


def _use_subprocess_embedder() -> bool:
    forced = tri_state_env("ADVISOR_EMBEDDING_SUBPROCESS")
    if forced is not None:
        return forced
    backend = os.environ.get("ADVISOR_MODEL_BACKEND", "").strip().lower()
    return platform.system() == "Darwin" and backend in {"minicpm", "minicpm-transformers"}


def _worker_loop() -> None:
    config_line = sys.stdin.readline()
    if not config_line:
        return
    embedder = LlamaCppEmbedder(**json.loads(config_line))
    for line in sys.stdin:
        if not line.strip():
            continue
        request = json.loads(line)
        request_id = request.get("id")
        try:
            vector = list(embedder.embed(str(request.get("text") or "")))
            response = {"id": request_id, "vector": vector}
        except Exception as error:
            response = {"id": request_id, "error": str(error)}
        print(json.dumps(response), flush=True)


if __name__ == "__main__":
    if len(sys.argv) == 2 and sys.argv[1] == "--worker":
        _worker_loop()
    else:
        raise SystemExit("usage: python -m hackathon_advisor.llama_embedding --worker")