dsh-biocore 0.1.0 → 0.1.2

package/python/biocore_memory/embedding_local.py

@@ -0,0 +1,447 @@
+# -*- coding: utf-8 -*-
+"""CPU-only local int8 ONNX embedder (the fallback path for remote embeddings).
+
+The optional component (task-190) is a bge-small-zh-v1.5 int8 ONNX export:
+
+  models/bge-small-zh-int8/model.onnx   about 23 MB
+  inputs   input_ids, attention_mask    int64, dynamic batch / sequence
+  output   embedding                    CLS-pooled + L2-normalized, dim 512
+
+Hard rules:
+- CPU only: onnxruntime with CPUExecutionProvider; no torch, no CUDA packages,
+  no GPU requirement; intra-op threads configurable (default 2).
+- Lazy import: importing this module never imports onnxruntime.
+- Clear failures: LocalEmbeddingUnavailable when runtime / model / vocab is
+  missing, LocalEmbeddingError for bad input or output dimensions. Callers
+  (the provider) catch these and stay on the remote path; the process never
+  dies because the optional component is absent.
+- Tokenizer: a small stdlib BERT WordPiece reader over vocab.txt, so the
+  component needs onnxruntime only (no transformers / tokenizers package).
+
+The model directory is an optional, removable component: delete it and the
+provider simply keeps using the remote API.
+"""
+from __future__ import annotations
+
+import json
+import os
+import threading
+import unicodedata
+from pathlib import Path
+
+DEFAULT_DIM = 512
+DEFAULT_MAX_LENGTH = 512
+DEFAULT_THREADS = 2
+MAX_THREADS = 8
+MAX_WORD_CHARS = 100
+SPECIAL_TOKENS = ("[PAD]", "[UNK]", "[CLS]", "[SEP]")
+ENV_TOKENIZER = "BIOCORE_EMBEDDING_LOCAL_TOKENIZER"
+
+# Lazy optional deps: imported on first real use, cached for the process, and
+# abstracted behind helpers so tests can inject fakes without touching
+# sys.modules (replacing sys.modules entries can force numpy to be loaded
+# twice in one process, which C extensions reject).
+_LAZY_MODULES = {}
+
+
+def _lazy_import(name):
+    module = _LAZY_MODULES.get(name)
+    if module is None:
+        module = __import__(name)
+        _LAZY_MODULES[name] = module
+    return module
+
+
+def _numpy():
+    return _lazy_import("numpy")
+
+
+def _onnxruntime():
+    return _lazy_import("onnxruntime")
+
+
+class LocalEmbeddingUnavailable(RuntimeError):
+    """The optional local component cannot be used (missing runtime/model)."""
+
+
+class LocalEmbeddingError(RuntimeError):
+    """The local component was used with bad input or produced bad output."""
+
+
+def _positive_int(value, default):
+    try:
+        number = int(value)
+    except Exception:
+        return int(default)
+    return number if number > 0 else int(default)
+
+
+def _chinese(char):
+    code = ord(char)
+    return (0x4E00 <= code <= 0x9FFF or 0x3400 <= code <= 0x4DBF
+            or 0x20000 <= code <= 0x2A6DF or 0x2A700 <= code <= 0x2B73F
+            or 0x2B740 <= code <= 0x2B81F or 0x2B820 <= code <= 0x2CEAF
+            or 0xF900 <= code <= 0xFAFF or 0x2F800 <= code <= 0x2FA1F)
+
+
+def _control(char):
+    category = unicodedata.category(char)
+    return category in ("Cc", "Cf") and char not in ("\t", "\n", "\r")
+
+
+def _whitespace(char):
+    if char in ("\t", "\n", "\r"):
+        return True
+    return unicodedata.category(char) in ("Zs", "Zl", "Zp")
+
+
+def _punctuation(char):
+    code = ord(char)
+    if (33 <= code <= 47 or 58 <= code <= 64 or 91 <= code <= 96
+            or 123 <= code <= 126):
+        return True
+    return unicodedata.category(char).startswith("P")
+
+
+def _clean_text(text, do_lower_case=True):
+    out = []
+    for char in str(text or ""):
+        code = ord(char)
+        if code == 0 or code == 0xFFFD:
+            continue
+        if _control(char):
+            continue
+        if _whitespace(char):
+            out.append(" ")
+            continue
+        out.append(char.lower() if do_lower_case else char)
+    return "".join(out)
+
+
+def _basic_tokens(text, do_lower_case=True):
+    """BERT BasicTokenizer: lowercase, CJK split, whitespace, punctuation."""
+    cleaned = _clean_text(text, do_lower_case)
+    padded = []
+    for char in cleaned:
+        if _chinese(char):
+            padded.append(" ")
+            padded.append(char)
+            padded.append(" ")
+        else:
+            padded.append(char)
+    tokens = []
+    for word in "".join(padded).split():
+        current = []
+        for char in word:
+            if _punctuation(char):
+                if current:
+                    tokens.append("".join(current))
+                    current = []
+                tokens.append(char)
+            else:
+                current.append(char)
+        if current:
+            tokens.append("".join(current))
+    return tokens
+
+
+def _wordpiece(token, vocab):
+    if len(token) > MAX_WORD_CHARS:
+        return [token]
+    chars = list(token)
+    pieces = []
+    start = 0
+    while start < len(chars):
+        end = len(chars)
+        found = None
+        while start < end:
+            piece = "".join(chars[start:end])
+            if start > 0:
+                piece = "##" + piece
+            if piece in vocab:
+                found = piece
+                break
+            end -= 1
+        if found is None:
+            return [token]
+        pieces.append(found)
+        start = end
+    return pieces
+
+
+class _VocabTokenizer:
+    """Minimal BERT WordPiece tokenizer over tokenizer.json / vocab.txt.
+
+    tokenizer.json (the fast-tokenizer export) is authoritative and is tried
+    first; vocab.txt is only a fallback for components that ship just the
+    classic file. Stdlib only.
+    """
+
+    def __init__(self, vocab_path, do_lower_case=True):
+        path = Path(vocab_path)
+        try:
+            if path.suffix.lower() == ".json":
+                payload = json.loads(
+                    path.read_text(encoding="utf-8", errors="replace"))
+                vocab = payload.get("model", {}).get("vocab")
+                if not isinstance(vocab, dict):
+                    raise ValueError("tokenizer.json has no model.vocab")
+                self.vocab = {}
+                for token, index in vocab.items():
+                    try:
+                        self.vocab[str(token)] = int(index)
+                    except (TypeError, ValueError):
+                        continue
+            else:
+                lines = path.read_text(
+                    encoding="utf-8", errors="replace").splitlines()
+                self.vocab = {}
+                for index, line in enumerate(lines):
+                    token = line[:-1] if line.endswith("\r") else line
+                    if token and token not in self.vocab:
+                        self.vocab[token] = index
+        except OSError as exc:
+            raise LocalEmbeddingUnavailable(
+                "cannot read tokenizer vocab: %s (%s)" % (path, exc)) from exc
+        except ValueError as exc:
+            raise LocalEmbeddingUnavailable(
+                "invalid tokenizer vocab %s (%s)" % (path, exc)) from exc
+        missing = [token for token in SPECIAL_TOKENS if token not in self.vocab]
+        if missing:
+            raise LocalEmbeddingUnavailable(
+                "tokenizer vocab is missing special tokens: %s"
+                % ", ".join(missing))
+        self.do_lower_case = bool(do_lower_case)
+        self.pad_id = self.vocab["[PAD]"]
+        self.unk_id = self.vocab["[UNK]"]
+        self.cls_id = self.vocab["[CLS]"]
+        self.sep_id = self.vocab["[SEP]"]
+
+    def encode(self, text, max_length=DEFAULT_MAX_LENGTH):
+        limit = max(8, _positive_int(max_length, DEFAULT_MAX_LENGTH))
+        pieces = []
+        for token in _basic_tokens(text, self.do_lower_case):
+            pieces.extend(_wordpiece(token, self.vocab))
+            if len(pieces) >= limit - 2:
+                break
+        pieces = pieces[:limit - 2]
+        ids = [self.cls_id]
+        ids.extend(self.vocab.get(piece, self.unk_id) for piece in pieces)
+        ids.append(self.sep_id)
+        return ids, [1] * len(ids)
+
+    def encode_batch(self, texts, max_length=DEFAULT_MAX_LENGTH):
+        pairs = [self.encode(text, max_length) for text in texts]
+        width = max([len(ids) for ids, _ in pairs] or [2])
+        ids_rows = []
+        mask_rows = []
+        for ids, mask in pairs:
+            padding = width - len(ids)
+            ids_rows.append(ids + [self.pad_id] * padding)
+            mask_rows.append(mask + [0] * padding)
+        return ids_rows, mask_rows
+
+
+def _as_texts(texts):
+    if texts is None:
+        return []
+    if isinstance(texts, str):
+        texts = [texts]
+    if not isinstance(texts, (list, tuple)):
+        raise LocalEmbeddingError(
+            "texts must be a string or a list of strings")
+    rows = []
+    for item in texts:
+        if item is None:
+            rows.append("")
+        elif isinstance(item, (dict, list, tuple, set)):
+            raise LocalEmbeddingError("text items must be strings")
+        else:
+            rows.append(str(item))
+    return rows
+
+
+def _normalize(values):
+    norm = 0.0
+    for value in values:
+        norm += value * value
+    norm = norm ** 0.5
+    if norm <= 1e-12:
+        return [0.0] * len(values)
+    return [value / norm for value in values]
+
+
+class LocalEmbedder:
+    """Lazy CPU-only int8 ONNX embedder. One instance may be shared."""
+
+    def __init__(self, model_path, dim=DEFAULT_DIM, tokenizer_path=None,
+                 max_length=DEFAULT_MAX_LENGTH, threads=DEFAULT_THREADS):
+        self.model_path = str(model_path or "").strip()
+        self.dim = _positive_int(dim, DEFAULT_DIM)
+        self.max_length = _positive_int(max_length, DEFAULT_MAX_LENGTH)
+        self.threads = max(1, min(MAX_THREADS,
+                                  _positive_int(threads, DEFAULT_THREADS)))
+        self.tokenizer_path = str(tokenizer_path or "").strip()
+        self._session = None
+        self._tokenizer = None
+        self._output_name = "embedding"
+        self._load_error = None
+        self._lock = threading.Lock()
+
+    # -- status -------------------------------------------------------
+    def status(self):
+        """ok | missing | failed (health string for /health)."""
+        if self._session is not None:
+            return "ok"
+        if self._load_error is not None:
+            return "failed"
+        if not self.model_path:
+            return "missing"
+        try:
+            self._model_file()
+        except LocalEmbeddingUnavailable:
+            return "failed"
+        return "missing"
+
+    def health(self):
+        return {"local_embedding": self.status(),
+                "model_path": self.model_path,
+                "dim": self.dim,
+                "threads": self.threads,
+                "load_error": str(self._load_error or "")[:200]}
+
+    # -- load ---------------------------------------------------------
+    def _model_file(self):
+        if not self.model_path:
+            raise LocalEmbeddingUnavailable(
+                "embedding_local_model_path is empty")
+        path = Path(self.model_path)
+        if path.is_dir():
+            path = path / "model.onnx"
+        if not path.is_file():
+            raise LocalEmbeddingUnavailable(
+                "local embedding model not found: %s" % path)
+        return path
+
+    def _vocab_file(self, model_file):
+        candidates = []
+
+        def add(path):
+            path = Path(path)
+            if path.suffix.lower() in (".json", ".txt"):
+                candidates.append(path)
+                return
+            candidates.append(path / "tokenizer.json")
+            candidates.append(path / "vocab.txt")
+
+        if self.tokenizer_path:
+            add(self.tokenizer_path)
+        env_path = os.environ.get(ENV_TOKENIZER, "").strip()
+        if env_path:
+            add(env_path)
+        add(model_file.parent)
+        add(model_file.parent / "tokenizer")
+        parent = model_file.parent.parent
+        if str(parent) != ".":
+            add(parent / "bge-small-zh")
+            add(parent / "bge-small-zh" / "tokenizer")
+        for candidate in candidates:
+            if candidate.is_file():
+                return candidate
+        raise LocalEmbeddingUnavailable(
+            "tokenizer (tokenizer.json or vocab.txt) not found for the local "
+            "embedding model; set embedding_local_tokenizer_path")
+
+    def load(self):
+        """Load onnxruntime + tokenizer once; raises LocalEmbeddingUnavailable."""
+        with self._lock:
+            if self._session is not None:
+                return self
+            if self._load_error is not None:
+                raise self._load_error
+            model_file = self._model_file()
+            vocab_file = self._vocab_file(model_file)
+            tokenizer = _VocabTokenizer(vocab_file)
+            try:
+                ort = _onnxruntime()
+            except Exception as exc:
+                error = LocalEmbeddingUnavailable(
+                    "onnxruntime is not installed; local embedding fallback "
+                    "is unavailable (%s)" % exc)
+                self._load_error = error
+                raise error
+            try:
+                options = ort.SessionOptions()
+                options.intra_op_num_threads = self.threads
+                options.inter_op_num_threads = 1
+                options.graph_optimization_level = (
+                    ort.GraphOptimizationLevel.ORT_ENABLE_ALL)
+                options.log_severity_level = 3
+                session = ort.InferenceSession(
+                    str(model_file), sess_options=options,
+                    providers=["CPUExecutionProvider"])
+            except Exception as exc:
+                error = LocalEmbeddingUnavailable(
+                    "failed to load local embedding model %s (%s)"
+                    % (model_file, exc))
+                self._load_error = error
+                raise error
+            try:
+                names = [item.name for item in session.get_outputs()]
+            except Exception:
+                names = []
+            if "embedding" in names or not names:
+                self._output_name = "embedding"
+            else:
+                self._output_name = names[0]
+            self._session = session
+            self._tokenizer = tokenizer
+            return self
+
+    # -- inference ----------------------------------------------------
+    def encode(self, texts):
+        """Encode a string or a list of strings into L2-normalized vectors."""
+        rows = _as_texts(texts)
+        if not rows:
+            return []
+        self.load()
+        ids, mask = self._tokenizer.encode_batch(rows, self.max_length)
+        try:
+            np = _numpy()
+        except Exception as exc:
+            raise LocalEmbeddingError(
+                "numpy is required for onnxruntime inputs (%s)" % exc)
+        feeds = {"input_ids": np.asarray(ids, dtype=np.int64),
+                 "attention_mask": np.asarray(mask, dtype=np.int64)}
+        try:
+            raw = self._session.run([self._output_name], feeds)[0]
+        except Exception as exc:
+            raise LocalEmbeddingError(
+                "local embedding inference failed (%s)" % exc)
+        output_rows = raw.tolist() if hasattr(raw, "tolist") else list(raw)
+        vectors = []
+        for row in output_rows:
+            try:
+                values = [float(item) for item in row]
+            except Exception as exc:
+                raise LocalEmbeddingError(
+                    "local embedding produced a non-numeric row (%s)" % exc)
+            if len(values) != self.dim:
+                raise LocalEmbeddingError(
+                    "dimension mismatch: expected %d, got %d"
+                    % (self.dim, len(values)))
+            vectors.append(_normalize(values))
+        return vectors
+
+    def warmup(self):
+        """Force load + one inference so startup can prove health."""
+        self.encode(["warmup"])
+        return self
+
+
+def create_local_embedder(model_path=None, dim=DEFAULT_DIM, tokenizer_path=None,
+                          threads=DEFAULT_THREADS, max_length=DEFAULT_MAX_LENGTH):
+    """Build an embedder without loading it (caller decides when to load)."""
+    return LocalEmbedder(model_path=model_path, dim=dim,
+                         tokenizer_path=tokenizer_path, threads=threads,
+                         max_length=max_length)

package/python/biocore_memory/embedding_mode.py

@@ -0,0 +1,156 @@
+# -*- coding: utf-8 -*-
+"""Embedding backend mode resolution (auto probe; read-only).
+
+task-211: embedding_mode now defaults to "auto" in the plugin schema. The
+selection is a pure, read-only probe:
+
+  a) remote keys complete (base_url + model + key) -> remote
+  b) else local model file exists and onnxruntime is importable -> local
+  c) else -> off (lexical/strength recall only; never an error)
+
+The probe never imports onnxruntime, never creates a session and never loads the
+model. The model is loaded lazily on the first real embed (see
+recall.embedding._LazyLocalEmbedder).
+
+Explicit values keep working unchanged: off / remote / local /
+remote_with_local_fallback. An old config without embedding_mode still uses the
+legacy loader path (requested == "legacy").
+"""
+from __future__ import annotations
+
+import importlib.util
+import os
+from pathlib import Path
+
+MODES = ("auto", "off", "remote", "local", "remote_with_local_fallback")
+REMOTE_KEYS = ("embedding_base_url", "embedding_model", "embedding_api_key")
+ENV_BASE_URL = "BIOCORE_EMBEDDING_BASE_URL"
+ENV_MODEL = "BIOCORE_EMBEDDING_MODEL"
+ENV_API_KEY = "BIOCORE_EMBEDDING_KEY"
+ENV_MODE = "BIOCORE_EMBEDDING_MODE"
+
+
+def _text(value):
+    return str(value or "").strip()
+
+
+def _flatten(config):
+    merged = {}
+    if isinstance(config, dict):
+        merged.update({key: value for key, value in config.items()
+                       if not isinstance(value, dict)})
+        for key in ("embedding", "embeddings", "provider"):
+            nested = config.get(key)
+            if isinstance(nested, dict):
+                merged.update(nested)
+    return merged
+
+
+def remote_settings(config=None):
+    """Remote endpoint keys with the same aliases/env fallbacks as recall."""
+    cfg = _flatten(config)
+    base_url = (cfg.get("embedding_base_url") or cfg.get("base_url")
+                or cfg.get("api_base") or cfg.get("url")
+                or cfg.get("embedding_api_base")
+                or os.environ.get(ENV_BASE_URL) or "")
+    model = (cfg.get("embedding_model") or cfg.get("model")
+             or os.environ.get(ENV_MODEL) or "")
+    api_key = (cfg.get("embedding_api_key") or cfg.get("api_key")
+               or cfg.get("key") or os.environ.get(ENV_API_KEY) or "")
+    return {"base_url": _text(base_url), "model": _text(model),
+            "api_key": _text(api_key)}
+
+
+def remote_complete(config=None):
+    """True only when base_url, model AND key are all non-empty."""
+    settings = remote_settings(config)
+    return bool(settings["base_url"] and settings["model"]
+                and settings["api_key"]), settings
+
+
+def local_model_file(model_path):
+    """Existing model.onnx path (a directory resolves to <dir>/model.onnx)."""
+    raw = _text(model_path)
+    if not raw:
+        return None
+    path = Path(raw)
+    if path.is_dir():
+        path = path / "model.onnx"
+    try:
+        return path if path.is_file() else None
+    except Exception:
+        return None
+
+
+def onnxruntime_available(find_spec=None):
+    """Importability probe only: never imports or initializes onnxruntime."""
+    try:
+        finder = find_spec if find_spec is not None else importlib.util.find_spec
+        return finder("onnxruntime") is not None
+    except Exception:
+        return False
+
+
+def resolve_mode(config=None, *, remote_ready=None, model_exists=None,
+                 onnx_available=None, find_spec=None):
+    """Resolve the effective embedding backend.
+
+    Injectable booleans make the probe testable without a filesystem/runtime.
+    Returns {requested, mode, source, reason, remote_ready,
+             local_model_exists, local_model_path, onnxruntime_available}.
+    """
+    cfg = _flatten(config)
+    raw = cfg.get("embedding_mode")
+    if raw is None:
+        raw = os.environ.get(ENV_MODE)
+    requested = _text(raw).lower()
+    if remote_ready is None:
+        remote_ready, _settings = remote_complete(cfg)
+    else:
+        remote_ready = bool(remote_ready)
+        _settings = remote_settings(cfg)
+    model_path = _text(cfg.get("embedding_local_model_path"))
+    if model_exists is None:
+        local_file = local_model_file(model_path)
+        local_model_exists = local_file is not None
+    else:
+        local_model_exists = bool(model_exists)
+        local_file = Path(model_path) if model_path else None
+    if onnx_available is None:
+        onnx_available = onnxruntime_available(find_spec)
+    else:
+        onnx_available = bool(onnx_available)
+    base = {
+        "requested": requested or "legacy",
+        "remote_ready": bool(remote_ready),
+        "local_model_exists": bool(local_model_exists),
+        "local_model_path": str(local_file or model_path or ""),
+        "onnxruntime_available": bool(onnx_available),
+    }
+    if not requested:
+        base.update({"mode": "", "source": "legacy",
+                     "reason": "旧配置未设置 embedding_mode：沿用旧规则"})
+        return base
+    if requested in ("off", "remote", "local", "remote_with_local_fallback"):
+        base.update({"mode": requested, "source": "explicit",
+                     "reason": "显式配置 embedding_mode=" + requested})
+        return base
+    if requested != "auto":
+        base.update({"mode": "off", "source": "explicit",
+                     "reason": "未知 embedding_mode=" + requested + "：按 off 处理"})
+        return base
+    if remote_ready:
+        base.update({"mode": "remote", "source": "auto",
+                     "reason": "auto→remote：远程配置齐全（base_url + model + key）"})
+        return base
+    if local_model_exists and onnx_available:
+        base.update({"mode": "local", "source": "auto",
+                     "reason": "auto→local：远程配置不全，但本地模型存在且 onnxruntime 可导入"})
+        return base
+    if local_model_exists and not onnx_available:
+        base.update({"mode": "off", "source": "auto",
+                     "reason": "auto→off：本地模型存在但 onnxruntime 不可导入"})
+        return base
+    base.update({"mode": "off", "source": "auto",
+                 "reason": "auto→off：远程配置不全且无本地模型"})
+    return base