build-small-hackathon/hackathon-advisor

1from __future__ import annotations2 3from collections.abc import Sequence4import atexit5import json6import os7from pathlib import Path8import platform9import subprocess10import sys11import threading12from typing import Any13 14from hackathon_advisor.config import bool_env, int_env, optional_int_env, tri_state_env15from hackathon_advisor.data import (16    DEFAULT_EMBEDDING_MODEL_FILE,17    DEFAULT_EMBEDDING_MODEL_REPO,18)19 20DEFAULT_N_CTX = 204821 22 23class LlamaCppEmbedder:24    def __init__(25        self,26        *,27        model_repo: str = DEFAULT_EMBEDDING_MODEL_REPO,28        model_file: str = DEFAULT_EMBEDDING_MODEL_FILE,29        model_path: str = "",30        n_ctx: int = DEFAULT_N_CTX,31        n_batch: int | None = None,32        n_threads: int | None = None,33        n_gpu_layers: int = 0,34        verbose: bool = False,35    ) -> None:36        self.model_repo = model_repo.strip() or DEFAULT_EMBEDDING_MODEL_REPO37        self.model_file = model_file.strip() or DEFAULT_EMBEDDING_MODEL_FILE38        self.model_path = model_path.strip()39        self.n_ctx = n_ctx40        self.n_batch = n_batch or n_ctx41        self.n_threads = n_threads42        self.n_gpu_layers = n_gpu_layers43        self.verbose = verbose44        self._model = None45 46    def __call__(self, text: str) -> Sequence[float]:47        return self.embed(text)48 49    def embed(self, text: str) -> Sequence[float]:50        model = self._ensure_model()51        return model.embed(text, normalize=True)52 53    def _ensure_model(self):54        if self._model is not None:55            return self._model56        from huggingface_hub import hf_hub_download57        from llama_cpp import LLAMA_POOLING_TYPE_MEAN, Llama58 59        model_path = self.model_path60        if not model_path:61            model_path = hf_hub_download(62                repo_id=self.model_repo,63                filename=self.model_file,64                repo_type="model",65            )66        if not Path(model_path).is_file():67            raise RuntimeError(f"llama.cpp embedding model was not found: {model_path}")68        self._model = Llama(69            model_path=model_path,70            embedding=True,71            pooling_type=LLAMA_POOLING_TYPE_MEAN,72            n_ctx=self.n_ctx,73            n_batch=self.n_batch,74            n_ubatch=self.n_batch,75            n_threads=self.n_threads,76            n_gpu_layers=self.n_gpu_layers,77            verbose=self.verbose,78        )79        return self._model80 81 82class SubprocessLlamaCppEmbedder:83    def __init__(84        self,85        *,86        model_repo: str = DEFAULT_EMBEDDING_MODEL_REPO,87        model_file: str = DEFAULT_EMBEDDING_MODEL_FILE,88        model_path: str = "",89        n_ctx: int = DEFAULT_N_CTX,90        n_batch: int | None = None,91        n_threads: int | None = None,92        n_gpu_layers: int = 0,93        verbose: bool = False,94    ) -> None:95        self.model_repo = model_repo.strip() or DEFAULT_EMBEDDING_MODEL_REPO96        self.model_file = model_file.strip() or DEFAULT_EMBEDDING_MODEL_FILE97        self.model_path = model_path.strip()98        self.n_ctx = n_ctx99        self.n_batch = n_batch or n_ctx100        self.n_threads = n_threads101        self.n_gpu_layers = n_gpu_layers102        self.verbose = verbose103        self._process: subprocess.Popen[str] | None = None104        self._request_id = 0105        self._lock = threading.Lock()106        atexit.register(self.close)107 108    def __call__(self, text: str) -> Sequence[float]:109        return self.embed(text)110 111    def embed(self, text: str) -> Sequence[float]:112        with self._lock:113            process = self._ensure_process()114            self._request_id += 1115            request_id = self._request_id116            request = json.dumps({"id": request_id, "text": text}, ensure_ascii=False)117            try:118                assert process.stdin is not None119                assert process.stdout is not None120                process.stdin.write(f"{request}\n")121                process.stdin.flush()122                line = process.stdout.readline()123            except (BrokenPipeError, OSError) as error:124                self.close()125                raise RuntimeError("llama.cpp embedding worker stopped before returning a vector.") from error126            if not line:127                returncode = process.poll()128                self.close()129                detail = f" with exit code {returncode}" if returncode is not None else ""130                raise RuntimeError(f"llama.cpp embedding worker exited{detail}.")131            try:132                response = json.loads(line)133            except json.JSONDecodeError as error:134                raise RuntimeError("llama.cpp embedding worker returned invalid JSON.") from error135            if response.get("id") != request_id:136                raise RuntimeError("llama.cpp embedding worker returned an out-of-order response.")137            if response.get("error"):138                raise RuntimeError(str(response["error"]))139            vector = response.get("vector")140            if not isinstance(vector, list):141                raise RuntimeError("llama.cpp embedding worker did not return a vector.")142            return vector143 144    def close(self) -> None:145        process = self._process146        self._process = None147        if process is None:148            return149        if process.poll() is None:150            process.terminate()151            try:152                process.wait(timeout=2)153            except subprocess.TimeoutExpired:154                process.kill()155                process.wait(timeout=2)156 157    def _ensure_process(self) -> subprocess.Popen[str]:158        if self._process is not None and self._process.poll() is None:159            return self._process160        self._process = subprocess.Popen(161            [sys.executable, "-u", "-m", "hackathon_advisor.llama_embedding", "--worker"],162            stdin=subprocess.PIPE,163            stdout=subprocess.PIPE,164            stderr=None if self.verbose else subprocess.DEVNULL,165            text=True,166            cwd=Path(__file__).resolve().parents[1],167        )168        config = json.dumps(169            {170                "model_repo": self.model_repo,171                "model_file": self.model_file,172                "model_path": self.model_path,173                "n_ctx": self.n_ctx,174                "n_batch": self.n_batch,175                "n_threads": self.n_threads,176                "n_gpu_layers": self.n_gpu_layers,177                "verbose": self.verbose,178            },179            ensure_ascii=False,180        )181        assert self._process.stdin is not None182        self._process.stdin.write(f"{config}\n")183        self._process.stdin.flush()184        return self._process185 186 187def create_llama_cpp_embedder(metadata: dict[str, Any]) -> LlamaCppEmbedder | SubprocessLlamaCppEmbedder:188    embedder_cls = SubprocessLlamaCppEmbedder if _use_subprocess_embedder() else LlamaCppEmbedder189    return embedder_cls(190        model_repo=os.environ.get(191            "ADVISOR_EMBEDDING_MODEL_REPO",192            str(metadata.get("model_repo") or DEFAULT_EMBEDDING_MODEL_REPO),193        ),194        model_file=os.environ.get(195            "ADVISOR_EMBEDDING_MODEL_FILE",196            str(metadata.get("model_file") or DEFAULT_EMBEDDING_MODEL_FILE),197        ),198        model_path=os.environ.get("ADVISOR_EMBEDDING_MODEL_PATH", ""),199        n_ctx=int_env("ADVISOR_EMBEDDING_N_CTX", DEFAULT_N_CTX, minimum=0),200        n_batch=optional_int_env("ADVISOR_EMBEDDING_BATCH"),201        n_threads=optional_int_env("ADVISOR_EMBEDDING_THREADS"),202        n_gpu_layers=int_env("ADVISOR_EMBEDDING_GPU_LAYERS", 0, minimum=0),203        verbose=bool_env("ADVISOR_EMBEDDING_VERBOSE"),204    )205 206 207def _use_subprocess_embedder() -> bool:208    forced = tri_state_env("ADVISOR_EMBEDDING_SUBPROCESS")209    if forced is not None:210        return forced211    backend = os.environ.get("ADVISOR_MODEL_BACKEND", "").strip().lower()212    return platform.system() == "Darwin" and backend in {"minicpm", "minicpm-transformers"}213 214 215def _worker_loop() -> None:216    config_line = sys.stdin.readline()217    if not config_line:218        return219    embedder = LlamaCppEmbedder(**json.loads(config_line))220    for line in sys.stdin:221        if not line.strip():222            continue223        request = json.loads(line)224        request_id = request.get("id")225        try:226            vector = list(embedder.embed(str(request.get("text") or "")))227            response = {"id": request_id, "vector": vector}228        except Exception as error:229            response = {"id": request_id, "error": str(error)}230        print(json.dumps(response), flush=True)231 232 233if __name__ == "__main__":234    if len(sys.argv) == 2 and sys.argv[1] == "--worker":235        _worker_loop()236    else:237        raise SystemExit("usage: python -m hackathon_advisor.llama_embedding --worker")238