陈世睿
917613c5ce
- crawler/:豆果(主)/下厨房/Bing 爬虫 + dishes.jsonl(3056 菜→9167 图映射)+ verify_repair 按 URL 重下 - retrieval/:三路检索(BM25 + 本地 BGE-M3 向量 + RRF 融合),FastAPI + 前端; 写死图片目录(默认 crawler/images,可 IMAGE_DIR 覆盖)、绑 0.0.0.0 局域网访问、 启动自动建索引、服务器 serve 图片 - 图片(1.7G)与向量模型(2.3G)不进 git Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
263 lines
9.1 KiB
Python
263 lines
9.1 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""图片名称检索测试台 —— 后端(FastAPI)。
|
|
|
|
三路检索,全部基于「图片名称」:
|
|
1) BM25 —— 关键词匹配(jieba 分词)
|
|
2) 向量 —— 本地 BAAI/bge-m3 余弦相似度
|
|
3) 融合 —— RRF 合并上面两路排名
|
|
|
|
【写死目录 + 局域网模式】图片目录写死在 IMAGE_DIR,服务器启动即扫描并自动建索引、
|
|
自己通过 /images 把图片发出去。绑 0.0.0.0,mentor 用局域网 IP 打开即可直接检索,
|
|
无需选文件夹、无需建索引。
|
|
|
|
跑法:python server.py → 控制台会打印可发给 mentor 的局域网地址。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import socket
|
|
import sys
|
|
import threading
|
|
|
|
import numpy as np
|
|
from fastapi import Body, FastAPI, Query
|
|
from fastapi.responses import JSONResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from bm25 import BM25Index
|
|
from embed_local import LocalEmbedder
|
|
from fuse import rrf
|
|
|
|
for _s in (sys.stdout, sys.stderr):
|
|
try:
|
|
_s.reconfigure(encoding="utf-8", errors="replace")
|
|
except Exception:
|
|
pass
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
# 图片目录:默认 = 本仓库的 crawler/images(爬虫默认就把图下到那里)。
|
|
# 想指向别处已爬好的图,启动前设环境变量 IMAGE_DIR=<绝对路径> 覆盖即可,无需改代码。
|
|
IMAGE_DIR = os.environ.get("IMAGE_DIR") or os.path.abspath(
|
|
os.path.join(HERE, os.pardir, "crawler", "images"))
|
|
IMG_EXT = (".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp")
|
|
PORT = 8799
|
|
|
|
app = FastAPI(title="图片名称检索测试台")
|
|
|
|
|
|
def scan_images(img_dir: str) -> dict[str, list[str]]:
|
|
"""扫描图片目录 → {菜名: [文件名,...]}。去掉尾部 _NN 把同名多图归并为一条。"""
|
|
name2files: dict[str, list[str]] = {}
|
|
if not os.path.isdir(img_dir):
|
|
return name2files
|
|
for fn in os.listdir(img_dir):
|
|
if not fn.lower().endswith(IMG_EXT):
|
|
continue
|
|
stem = os.path.splitext(fn)[0]
|
|
name = re.sub(r"_\d+$", "", stem) or stem
|
|
name2files.setdefault(name, []).append(fn)
|
|
for files in name2files.values():
|
|
files.sort()
|
|
return name2files
|
|
|
|
|
|
class State:
|
|
"""全局索引状态(单实例工具,够用)。"""
|
|
|
|
def __init__(self):
|
|
self.lock = threading.Lock()
|
|
self.names: list[str] = []
|
|
self.name2files: dict[str, list[str]] = {}
|
|
self.bm25: BM25Index | None = None
|
|
self.embedder: LocalEmbedder | None = None
|
|
self.model = "BAAI/bge-m3"
|
|
self.vec_names: list[str] = []
|
|
self.vec_matrix: np.ndarray | None = None
|
|
self.state = "idle" # idle | building | ready | error
|
|
self.phase = "" # bm25 | loadmodel | embedding | done
|
|
self.emb_done = 0
|
|
self.emb_total = 0
|
|
self.failed = 0
|
|
self.message = ""
|
|
self.device = ""
|
|
|
|
|
|
ST = State()
|
|
|
|
|
|
def build_index(names: list[str], model: str) -> None:
|
|
try:
|
|
with ST.lock:
|
|
ST.state, ST.phase, ST.message = "building", "bm25", ""
|
|
ST.names, ST.model = names, model
|
|
ST.bm25 = None
|
|
ST.vec_matrix, ST.vec_names = None, []
|
|
ST.emb_done = ST.emb_total = ST.failed = 0
|
|
|
|
# 1) BM25(秒级)
|
|
bm = BM25Index(names)
|
|
with ST.lock:
|
|
ST.bm25, ST.phase = bm, "loadmodel"
|
|
|
|
# 2) 本地 BGE-M3:先加载模型(首次会下载约 2.3G),再编码(有 CUDA 走 GPU)
|
|
emb = LocalEmbedder(model=model, cache_dir=os.path.join(HERE, "cache"))
|
|
dev = emb.load_model()
|
|
with ST.lock:
|
|
ST.device, ST.phase = dev, "embedding"
|
|
|
|
def prog(done: int, total: int) -> None:
|
|
with ST.lock:
|
|
ST.emb_done, ST.emb_total = done, total
|
|
|
|
emb.embed_corpus(names, progress=prog)
|
|
vnames, mat = emb.matrix(names)
|
|
with ST.lock:
|
|
ST.embedder, ST.vec_names, ST.vec_matrix = emb, vnames, mat
|
|
ST.failed = emb.last_failed
|
|
ST.message = "" if vnames else ("向量未生成:" + emb.last_error if emb.last_error else "")
|
|
ST.state, ST.phase = "ready", "done"
|
|
except Exception as e: # noqa: BLE001
|
|
with ST.lock:
|
|
ST.state, ST.message = "error", f"{type(e).__name__}: {e}"
|
|
|
|
|
|
@app.on_event("startup")
|
|
def _startup_autobuild() -> None:
|
|
"""启动即扫描写死目录并后台建索引,mentor 打开就能搜。"""
|
|
n2f = scan_images(IMAGE_DIR)
|
|
with ST.lock:
|
|
ST.name2files = n2f
|
|
names = sorted(n2f.keys())
|
|
if not names:
|
|
with ST.lock:
|
|
ST.state = "error"
|
|
ST.message = f"图片目录为空或不存在:{IMAGE_DIR}"
|
|
return
|
|
threading.Thread(target=build_index, args=(names, "BAAI/bge-m3"), daemon=True).start()
|
|
|
|
|
|
@app.post("/api/index")
|
|
def api_index(payload: dict = Body(...)):
|
|
"""(保留)手动指定名称重建索引;写死目录模式下一般用不到。"""
|
|
names = [str(n).strip() for n in payload.get("names", []) if str(n).strip()]
|
|
names = list(dict.fromkeys(names))
|
|
model = (payload.get("model") or "BAAI/bge-m3").strip()
|
|
if not names:
|
|
return JSONResponse({"ok": False, "error": "没有菜名"}, status_code=400)
|
|
if ST.state == "building":
|
|
return JSONResponse({"ok": False, "error": "正在建索引,请稍候"}, status_code=409)
|
|
threading.Thread(target=build_index, args=(names, model), daemon=True).start()
|
|
return {"ok": True, "total": len(names), "model": model}
|
|
|
|
|
|
@app.get("/api/status")
|
|
def api_status():
|
|
with ST.lock:
|
|
return {
|
|
"state": ST.state,
|
|
"phase": ST.phase,
|
|
"bm25_ready": ST.bm25 is not None,
|
|
"vec_ready": ST.vec_matrix is not None and len(ST.vec_names) > 0,
|
|
"total": len(ST.names),
|
|
"n_images": sum(len(v) for v in ST.name2files.values()),
|
|
"emb_done": ST.emb_done,
|
|
"emb_total": ST.emb_total,
|
|
"vec_count": len(ST.vec_names),
|
|
"failed": ST.failed,
|
|
"model": ST.model,
|
|
"device": ST.device,
|
|
"message": ST.message,
|
|
}
|
|
|
|
|
|
def _bm25_search(q: str, topk: int) -> list[tuple[str, float]]:
|
|
return ST.bm25.search(q, topk=topk) if ST.bm25 else []
|
|
|
|
|
|
def _vec_search(q: str, topk: int) -> list[tuple[str, float]]:
|
|
mat, names, emb = ST.vec_matrix, ST.vec_names, ST.embedder
|
|
if mat is None or len(names) == 0 or emb is None:
|
|
return []
|
|
qv = emb.embed_query(q)
|
|
sims = mat @ qv
|
|
order = np.argsort(sims)[::-1][:topk]
|
|
return [(names[i], float(sims[i])) for i in order]
|
|
|
|
|
|
@app.get("/api/search")
|
|
def api_search(q: str = Query(...), topk: int = 30, rrf_k: int = 60,
|
|
w_bm25: float = 1.0, w_vec: float = 1.0):
|
|
q = q.strip()
|
|
if not q:
|
|
return {"q": q, "bm25": [], "vector": [], "fusion": []}
|
|
|
|
bm = _bm25_search(q, topk)
|
|
ve = _vec_search(q, topk)
|
|
bm_rank = [n for n, _ in bm]
|
|
ve_rank = [n for n, _ in ve]
|
|
fused = rrf({"bm25": bm_rank, "vector": ve_rank},
|
|
{"bm25": w_bm25, "vector": w_vec}, k=rrf_k)[:topk]
|
|
|
|
bm_s, ve_s = dict(bm), dict(ve)
|
|
bm_pos = {n: i + 1 for i, n in enumerate(bm_rank)}
|
|
ve_pos = {n: i + 1 for i, n in enumerate(ve_rank)}
|
|
|
|
def card(name: str, score: float) -> dict:
|
|
return {
|
|
"name": name,
|
|
"score": round(score, 4),
|
|
"files": ST.name2files.get(name, []), # 服务器 serve 的图片文件名
|
|
"bm25": round(bm_s[name], 3) if name in bm_s else None,
|
|
"bm25_rank": bm_pos.get(name),
|
|
"vec": round(ve_s[name], 4) if name in ve_s else None,
|
|
"vec_rank": ve_pos.get(name),
|
|
}
|
|
|
|
return {
|
|
"q": q,
|
|
"bm25": [card(n, s) for n, s in bm],
|
|
"vector": [card(n, s) for n, s in ve],
|
|
"fusion": [card(n, s) for n, s in fused],
|
|
}
|
|
|
|
|
|
# 图片由服务器 serve(局域网里 mentor 的浏览器据此显示);必须在 "/" 之前挂载。
|
|
if os.path.isdir(IMAGE_DIR):
|
|
app.mount("/images", StaticFiles(directory=IMAGE_DIR), name="images")
|
|
# 静态前端(放在所有 /api 与 /images 之后挂载)
|
|
app.mount("/", StaticFiles(directory=os.path.join(HERE, "static"), html=True), name="static")
|
|
|
|
|
|
def _lan_ips() -> list[str]:
|
|
ips: list[str] = []
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
s.connect(("8.8.8.8", 80))
|
|
ips.append(s.getsockname()[0])
|
|
s.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
|
|
ip = info[4][0]
|
|
if ip not in ips:
|
|
ips.append(ip)
|
|
except Exception:
|
|
pass
|
|
return [i for i in ips if not i.startswith("127.")]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
|
|
print("=" * 56, flush=True)
|
|
print(f"图片目录:{IMAGE_DIR}", flush=True)
|
|
print("发给 mentor 的局域网地址:", flush=True)
|
|
for ip in _lan_ips() or ["<查不到,用 ipconfig 看本机 IPv4>"]:
|
|
print(f" http://{ip}:{PORT}", flush=True)
|
|
print(f"本机自测:http://127.0.0.1:{PORT}", flush=True)
|
|
print("(若 mentor 连不上:Windows 防火墙放行该端口入站)", flush=True)
|
|
print("=" * 56, flush=True)
|
|
uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")
|