tzt
f783878c02
- embedder.py：embed() 调 llama-server /v1/embeddings -> 对称 per-vector int8 量化（scale=127/max|v|，1024 维余弦扰动 ~1e-4 << 1e-2 验收）； 超时/连接/非200/空形状 -> EmbedderDown（D-G4 fail-closed）； quantize_int8/cosine_int8 辅助（int8 直接点积，scale 正标量不改方向） - routes.py 拆双路由组：/sense 前缀组 + /v1 无前缀组（/v1/embeddings OpenAI 兼容透传，embedder 不可用 503）；build_sense_routers 返回列表， api.py 逐个 include；build_sense_router 保留向后兼容 - 测试 +6：量化余弦误差 200 组 <1e-2/范围/零向量/int8 余弦/ 正常量化/三种 fail-closed/透传端点（503+形状+400），全量 381 passed - 待真机项：llama-server embedder 端点冒烟（需用户配置 embedder.base_url）
85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
"""Embedder 客户端（T-G1）：llama-server /v1/embeddings + int8 量化 + 降级（D-G4）。
|
||
|
||
- int8 量化：对称 per-vector scale（scale = max|v| / 127），q = round(v × scale)。
|
||
1024 维余弦扰动 ~1e-4（验收 <1e-2）；存储 1024 维 fp32 4KB -> int8 1KB（D-G6）。
|
||
- 降级（D-G4）：超时/连接失败/非 200/形状不符 -> raise EmbedderDown，
|
||
调用方（grader）退化为规则门 + 默认 T2。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import httpx
|
||
|
||
from gateway.sense.config import EmbedderCfg
|
||
from gateway.sense.errors import EmbedderDown
|
||
|
||
_client: Optional[httpx.AsyncClient] = None
|
||
|
||
|
||
def get_client() -> httpx.AsyncClient:
|
||
global _client
|
||
if _client is None:
|
||
_client = httpx.AsyncClient(timeout=5.0)
|
||
return _client
|
||
|
||
|
||
async def close_client() -> None:
|
||
global _client
|
||
if _client is not None:
|
||
await _client.aclose()
|
||
_client = None
|
||
|
||
|
||
def quantize_int8(vec: list[float]) -> list[int]:
|
||
"""对称 per-vector int8 量化：q = round(v × scale)，scale = 127 / max|v|。"""
|
||
m = max((abs(v) for v in vec), default=0.0)
|
||
if m == 0:
|
||
return [0] * len(vec)
|
||
scale = 127.0 / m
|
||
return [max(-127, min(127, int(round(v * scale)))) for v in vec]
|
||
|
||
|
||
def dequantize_int8(q: list[int], scale: float) -> list[float]:
|
||
"""反量化（需要保留 scale 时使用；检索场景直接用 int8 余弦即可）。"""
|
||
if scale == 0:
|
||
return [0.0] * len(q)
|
||
return [x / scale for x in q]
|
||
|
||
|
||
def cosine_int8(a: list[int], b: list[int]) -> float:
|
||
"""int8 向量余弦相似度（无 scale 版本：scale 为正标量不改变方向）。"""
|
||
n = min(len(a), len(b))
|
||
if n == 0:
|
||
return 0.0
|
||
dot = na = nb = 0
|
||
for i in range(n):
|
||
x, y = a[i], b[i]
|
||
dot += x * y
|
||
na += x * x
|
||
nb += y * y
|
||
if na == 0 or nb == 0:
|
||
return 0.0
|
||
return max(-1.0, min(1.0, dot / ((na ** 0.5) * (nb ** 0.5))))
|
||
|
||
|
||
async def embed(text: str, cfg: EmbedderCfg) -> list[int]:
|
||
"""调用 llama-server /v1/embeddings，返回 int8 量化向量（§7 签名）。
|
||
|
||
失败（超时/连接/非 200/形状不符）-> raise EmbedderDown（D-G4 降级信号）。
|
||
"""
|
||
try:
|
||
client = get_client()
|
||
resp = await client.post(
|
||
cfg.base_url.rstrip("/") + "/embeddings",
|
||
json={"model": cfg.model, "input": text},
|
||
timeout=cfg.timeout_s)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
vec = data["data"][0]["embedding"]
|
||
if not isinstance(vec, list) or not vec:
|
||
raise ValueError("空向量")
|
||
return quantize_int8([float(v) for v in vec])
|
||
except EmbedderDown:
|
||
raise
|
||
except Exception as e: # noqa: BLE001
|
||
raise EmbedderDown(f"embedder 不可用: {type(e).__name__}: {e}") from e
|