File size: 7,681 Bytes
750b88d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Model access for the graph-memory engines (HippoRAG 2 and ReMem).
The engines call a language model through an object with `infer(messages, **kw)` returning
`(text, metadata, cache_hit)` and embed text through an object with `encode` / `batch_encode`.
`EngineChat` and `EngineEmbedder` provide those interfaces on top of the evaluation's chat client
and an OpenAI-compatible embeddings endpoint, so the engines use the configured backbone and the
token usage of every call is logged under the current phase.
"""
from __future__ import annotations
import json
import logging
import os
import sys
import time
import types
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import numpy as np
from ..llm_client import ENDPOINTS, MAX_ATTEMPTS, SDK_MAX_RETRIES, USAGE, api_key_for, http_timeout
logger = logging.getLogger(__name__)
VENDOR = Path(__file__).resolve().parents[2] / "vendor"
_FENCE_OPEN, _FENCE_CLOSE = "```json\n", "```"
def add_engine_source(name: str) -> None:
"""Put `vendor/<name>/src` on the import path."""
src = VENDOR / name / "src"
if not src.exists():
raise ImportError(f"{name} source folder not found at {src}")
if str(src) not in sys.path:
sys.path.insert(0, str(src))
def stub_optional_imports() -> None:
"""HippoRAG imports vllm at module level for its offline modes, which are not used here.
When vllm is absent, register empty modules carrying the names those imports expect."""
if "vllm" in sys.modules:
return
try:
import vllm # noqa: F401
return
except ImportError:
pass
names = {"vllm": ("SamplingParams", "LLM"),
"vllm.model_executor": (),
"vllm.model_executor.guided_decoding": (),
"vllm.model_executor.guided_decoding.guided_fields": ("GuidedDecodingRequest",)}
for module_name, classes in names.items():
module = types.ModuleType(module_name)
for cls in classes:
setattr(module, cls, type(cls, (), {}))
sys.modules[module_name] = module
logger.info("vllm is not installed; empty placeholder modules registered")
def read_openie_docs(path: str) -> list:
"""The `docs` list of an engine's OpenIE cache file; empty when the file is missing or unreadable."""
if not os.path.exists(path):
return []
try:
with open(path, encoding="utf-8") as fh:
return json.load(fh).get("docs", [])
except Exception as e:
logger.warning("could not read the OpenIE cache %s: %s", path, e)
return []
class EngineChat:
"""`infer` and `batch_infer` over the evaluation's chat client.
With `max_tokens` set, each call sends an output limit: the engine's own `max_completion_tokens`
or `max_tokens` when it passes one, else `max_tokens`. With `max_tokens=None` no limit is sent
and the client's default applies. `echo=True` adds the messages and the reply to the metadata.
The metadata carries the reply's finish reason ("stop" when the client reports none), so an
engine can repair JSON that was cut off at the output limit ("length").
"""
def __init__(self, client, *, temperature: float, max_tokens: int | None = None, echo: bool = False):
self._client = client
self._temperature = temperature
self._max_tokens = max_tokens
self._echo = echo
def infer(self, messages, **kw):
request = {"temperature": kw.get("temperature", self._temperature)}
if self._max_tokens is not None:
request["max_tokens"] = kw.get("max_completion_tokens", kw.get("max_tokens", self._max_tokens))
if "response_format" in kw:
request["response_format"] = kw["response_format"]
try:
reply = self._client.chat(messages=messages, **request)
text = reply.content
wants_json = (kw.get("response_format") or {}).get("type") == "json_object"
if wants_json and text.startswith(_FENCE_OPEN) and text.endswith(_FENCE_CLOSE):
text = text[len(_FENCE_OPEN):-len(_FENCE_CLOSE)].strip()
except Exception as e:
logger.error("engine chat call failed: %s", e)
raise
meta = {"prompt": messages, "response": text} if self._echo else {}
meta.update(prompt_tokens=reply.input_tokens, completion_tokens=reply.output_tokens,
finish_reason=getattr(reply, "finish_reason", None) or "stop")
return text, meta, False
def batch_infer(self, messages_list, max_workers: int = 10, **kw):
"""`infer` over several message lists in parallel; results keep the input order."""
with ThreadPoolExecutor(max_workers=max_workers) as pool:
pending = [pool.submit(self.infer, messages, **kw) for messages in messages_list]
return [p.result() for p in pending]
class EngineEmbedder:
"""`encode` / `batch_encode` over an OpenAI-compatible embeddings endpoint.
Blank texts are sent as "empty"; an `instruction` keyword is prefixed to every text. Requests
hold at most `batch_size` texts and are retried up to MAX_ATTEMPTS times (a 400 response is
not retried). Rows are scaled to unit length unless `norm=False` is passed.
"""
def __init__(self, *, model: str | None, batch_size: int | None, dim: int | None = None,
api_key: str | None = None, base_url: str | None = None):
self.embedding_model_name = model
self.embedding_size = dim
self._batch_size = batch_size
self._api_key = api_key
self._base_url = base_url
self._endpoint = None
def _embeddings(self):
if self._endpoint is None:
from openai import OpenAI
base_url = self._base_url or ENDPOINTS["embedding"]
self._endpoint = OpenAI(api_key=self._api_key or api_key_for(base_url), base_url=base_url,
timeout=http_timeout(), max_retries=SDK_MAX_RETRIES).embeddings
logger.info("embedding endpoint ready, model=%s", self.embedding_model_name)
return self._endpoint
def _request(self, texts: list) -> list:
attempt = 0
while True:
attempt += 1
try:
response = self._embeddings().create(input=texts, model=self.embedding_model_name)
except Exception as e:
if attempt < MAX_ATTEMPTS and "Error code: 400" not in str(e):
time.sleep(1.5 * attempt)
continue
raise RuntimeError(f"embedding request failed after {attempt} attempt(s): {e}") from e
usage = getattr(response, "usage", None)
if usage is not None:
USAGE.log(self.embedding_model_name, "embed", int(getattr(usage, "prompt_tokens", 0) or 0), 0)
return [row.embedding for row in response.data]
def encode(self, texts, **kw) -> np.ndarray:
texts = [texts] if isinstance(texts, str) else texts
prefix = kw.get("instruction", "")
texts = [f"{prefix}{t if t and t.strip() else 'empty'}" for t in texts]
if not self._batch_size:
raise RuntimeError("set embedding.batch_size in backbone.yaml")
step = int(self._batch_size)
vectors = []
for start in range(0, len(texts), step):
vectors.extend(self._request(texts[start:start + step]))
matrix = np.array(vectors)
if kw.get("norm", True):
lengths = np.linalg.norm(matrix, axis=1, keepdims=True)
matrix = matrix / np.where(lengths == 0, 1, lengths)
return matrix
batch_encode = encode
|