From dbb9a032aabd1c893af929250678eafa5f9fb376 Mon Sep 17 00:00:00 2001 From: xiaowei Date: Sat, 30 May 2026 03:18:43 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E:=20=E6=9C=AC=E5=9C=B0=20bge-?= =?UTF-8?q?m3=20=E5=B5=8C=E5=85=A5=E6=9C=8D=E5=8A=A1=E5=99=A8=20(embed-ser?= =?UTF-8?q?ver.py)=20+=20=E9=87=8D=E7=B4=A2=E5=BC=95=E8=84=9A=E6=9C=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- deploy/embed-server.py | 84 ++++++++++++++++++++++ scripts/reindex_bge_m3.py | 147 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 231 insertions(+) create mode 100644 deploy/embed-server.py create mode 100644 scripts/reindex_bge_m3.py diff --git a/deploy/embed-server.py b/deploy/embed-server.py new file mode 100644 index 0000000..e4a64e6 --- /dev/null +++ b/deploy/embed-server.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""轻量 bge-m3 embedding server — OpenAI 兼容 /v1/embeddings 接口""" +import os, sys, json, time +from typing import List + +# 环境变量 +HOST = os.environ.get("EMBED_HOST", "0.0.0.0") +PORT = int(os.environ.get("EMBED_PORT", "8000")) +MODEL = os.environ.get("EMBED_MODEL", "BAAI/bge-m3") +DEVICE = os.environ.get("EMBED_DEVICE", "cpu") + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse +from pydantic import BaseModel, Field +import uvicorn + +app = FastAPI(title="bge-m3 Embedding Server", version="1.0") + +class EmbedRequest(BaseModel): + input: str | List[str] + model: str = MODEL + +class EmbeddingObject(BaseModel): + object: str = "embedding" + index: int + embedding: List[float] + +class UsageInfo(BaseModel): + prompt_tokens: int + total_tokens: int + +class EmbedResponse(BaseModel): + object: str = "list" + data: List[EmbeddingObject] + model: str = MODEL + usage: UsageInfo + +# 延迟加载 +model_pipe = None + +def load_model(): + global model_pipe + if model_pipe is not None: + return + print(f"[embed] Loading {MODEL} on {DEVICE}...", flush=True) + t0 = time.time() + from sentence_transformers import SentenceTransformer + model_pipe = SentenceTransformer(MODEL, device=DEVICE) + elapsed = time.time() - t0 + print(f"[embed] Loaded in {elapsed:.1f}s. Dimension: {model_pipe.get_sentence_embedding_dimension()}", flush=True) + +@app.on_event("startup") +async def startup(): + load_model() + +@app.get("/health") +async def health(): + return {"status": "ok", "model": MODEL, "device": DEVICE, + "loaded": model_pipe is not None} + +@app.post("/v1/embeddings") +async def embed(req: EmbedRequest): + t0 = time.time() + texts = req.input if isinstance(req.input, list) else [req.input] + + # bge-m3 需要加 prefix + prefixed = [f"为这个句子生成表示以用于检索相关文章：{t}" for t in texts] + + vecs = model_pipe.encode(prefixed, normalize_embeddings=True, + show_progress_bar=False) + vecs = vecs.tolist() + + total_tokens = sum(len(t) for t in texts) + elapsed = time.time() - t0 + + data = [ + EmbeddingObject(index=i, embedding=vecs[i]) + for i in range(len(vecs)) + ] + return EmbedResponse(data=data, usage=UsageInfo(prompt_tokens=total_tokens, total_tokens=total_tokens)) + +if __name__ == "__main__": + load_model() + uvicorn.run(app, host=HOST, port=PORT, log_level="info") diff --git a/scripts/reindex_bge_m3.py b/scripts/reindex_bge_m3.py new file mode 100644 index 0000000..089a0df --- /dev/null +++ b/scripts/reindex_bge_m3.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +"""存量记忆向量重编码：将 LanceDB 中所有 memory vector 替换为本地 bge-m3 编码 + +策略：全量读 → 分批编码 → 删表重建（一次写入，1130 条 batch 插入） +""" + +import os, sys, json, time, argparse +from typing import List + +import lancedb +import pyarrow as pa +import requests +import numpy as np + +EMBED_URL = os.environ.get("EMBED_URL", "http://localhost:8000/v1/embeddings") +LANCE_DIR = os.environ.get("LANCE_DIR", "/var/lib/memoryweave/lancedb") +ENCODE_BATCH = 3 # 每批编码 3 条，避开 sentence-transformers batch bug + +def encode_batch(texts: List[str]) -> List[List[float]]: + """调用本地 bge-m3 server 批量编码（server 自带 prefix）""" + resp = requests.post(EMBED_URL, json={"input": texts, "model": "BAAI/bge-m3"}, + timeout=60) + resp.raise_for_status() + data = resp.json()["data"] + vecs = [None] * len(texts) + for item in data: + vecs[item["index"]] = item["embedding"] + return vecs + +def main(): + parser = argparse.ArgumentParser(description="存量记忆向量重编码") + parser.add_argument("--backup", action="store_true", help="备份当前 LanceDB") + parser.add_argument("--dry-run", action="store_true", help="只读+编码，不写入") + parser.add_argument("--batch", type=int, default=ENCODE_BATCH, help="编码批次大小") + parser.add_argument("--embed-url", default=EMBED_URL) + args = parser.parse_args() + + # 1. 验证嵌入服务 + try: + r = requests.get(args.embed_url.replace("/v1/embeddings", "/health"), timeout=5) + r.raise_for_status() + status = r.json() + print(f"[health] model={status.get('model','?')} device={status.get('device','?')}") + except Exception as e: + print(f"[FATAL] 嵌入服务不可用: {e}") + sys.exit(1) + + # 2. 打开 LanceDB，读取全量数据 + db = lancedb.connect(LANCE_DIR) + tbl = db.open_table("memories") + arrow = tbl.to_arrow() + total = arrow.num_rows + print(f"[lance] 读取 {total} 条记录") + + col_names = arrow.schema.names + print(f"[schema] 列: {col_names}") + + # 3. 备份 + if args.backup and not args.dry_run: + import shutil + bak_dir = f"{LANCE_DIR}.bak.{int(time.time())}" + print(f"[backup] 备份到 {bak_dir} ...") + shutil.copytree(LANCE_DIR, bak_dir) + print(f"[backup] 完成") + + # 4. 读取 content + ids = arrow.column("id").to_pylist() + contents = arrow.column("content").to_pylist() + print(f"[data] 首条: {repr(contents[0][:60]) if contents else 'EMPTY'}") + + # 5. 分批编码 + print(f"[encode] 分批编码 {total} 条 (batch_size={args.batch}) ...") + all_vecs = [None] * total + errors = 0 + t0 = time.time() + + for offset in range(0, total, args.batch): + batch_texts = contents[offset:offset + args.batch] + try: + vecs = encode_batch(batch_texts) + for i, v in enumerate(vecs): + all_vecs[offset + i] = v + except Exception as e: + # 单条回退 + for i in range(len(batch_texts)): + try: + v = encode_batch([batch_texts[i]]) + all_vecs[offset + i] = v[0] + except Exception as e2: + print(f" [ERROR] id={ids[offset+i]}: {e2}") + errors += 1 + + elapsed = time.time() - t0 + done = min(offset + args.batch, total) + rate = done / elapsed if elapsed > 0 else 0 + print(f" {done}/{total} ({rate:.1f}/s) errors={errors}", end="\r") + + print() + elapsed = time.time() - t0 + + success = total - errors + print(f"[encode] {elapsed:.1f}s | {success} 成功, {errors} 失败") + + if errors > 0: + # 用零向量填充失败项 + print("[encode] 用零向量填充失败项") + zero = [0.0] * 1024 + for i in range(total): + if all_vecs[i] is None: + all_vecs[i] = zero + + if args.dry_run: + print("[dry-run] 完成，未写入 LanceDB") + return + + # 6. 建新 vector 列（PyArrow FixedShapeTensorArray from numpy） + vec_array = np.array(all_vecs, dtype=np.float32) # (total, 1024) + vec_col = pa.FixedShapeTensorArray.from_numpy_ndarray(vec_array) + + # 替换原有 vector 列 + col_index = col_names.index("vector") + new_columns = [] + for i, name in enumerate(col_names): + if name == "vector": + new_columns.append(vec_col) + else: + new_columns.append(arrow.column(name)) + + new_arrow = pa.table(dict(zip(col_names, new_columns))) + + # 7. 删表重建 + print("[write] 删除旧表 ...") + db.drop_table("memories") + print("[write] 重建 memories 表 ...") + db.create_table("memories", new_arrow) + print("[write] 完成 ✓") + + # 8. 验证 + verify_tbl = db.open_table("memories") + verify_count = verify_tbl.count_rows() + verify_vec = verify_tbl.to_arrow().column("vector")[0] + verify_dim = len(verify_vec) if verify_vec else 0 + print(f"[verify] 新表: {verify_count} 条, 向量维度={verify_dim}") + print(f"[done] 总耗时 {time.time()-t0:.1f}s") + +if __name__ == "__main__": + main()