| from FlagEmbedding import BGEM3FlagModel, FlagReranker |
| import numpy as np |
| import json |
| from pathlib import Path |
| import faiss |
| import soundfile as sf |
|
|
| def load_jsonl(path): |
| out = [] |
| with open(path, "r", encoding="utf-8") as f: |
| for ln in f: |
| ln = ln.strip() |
| if ln: |
| out.append(json.loads(ln)) |
| return out |
|
|
| def load_vector_db(faiss_index_path, chunks_path, meta_path, uid_list_path): |
| print("Loading vector database...") |
| index = faiss.read_index(faiss_index_path) |
| chunks = load_jsonl(chunks_path) |
| meta = load_jsonl(meta_path) |
|
|
| print(f"FAISS ntotal: {index.ntotal}, Chunks: {len(chunks)}, Meta: {len(meta)}") |
| assert index.ntotal == len(chunks) == len(meta), "DB mapping mismatch!" |
| return index, chunks, meta |
|
|
| def load_embedding_models(emb_model_dir, rerank_model_dir): |
| print("Loading embedding and reranking models...") |
| embedder = BGEM3FlagModel(emb_model_dir, use_fp16=True) |
| reranker = FlagReranker(rerank_model_dir, use_fp16=True) |
| print("Embedding and reranking models loaded successfully!") |
| return embedder, reranker |
|
|
| def l2_normalize(v): |
| denom = np.linalg.norm(v, axis=1, keepdims=True) + 1e-12 |
| return v / denom |
|
|
| def embed_query(embedder, q): |
| out = embedder.encode([q], return_dense=True, return_sparse=False, return_colbert_vecs=False) |
| vec = np.array(out["dense_vecs"], dtype=np.float32) |
| return l2_normalize(vec).astype("float32") |
|
|
| def retrieve_then_rerank(embedder, reranker, index, chunks, meta, question_bn, top_k=3, faiss_top_n=30): |
| """RAG: Retrieve from FAISS and rerank""" |
| q_for_retrieval = question_bn |
| |
| qvec = embed_query(embedder, q_for_retrieval) |
| D, I = index.search(qvec, faiss_top_n) |
| cand_idxs = I[0].tolist() |
| cand_texts = [] |
| valid_idxs = [] |
| |
| for i in cand_idxs: |
| t = chunks[i].get("text") or chunks[i].get("chunk_text") or "" |
| t = (t or "").strip() |
| if t: |
| cand_texts.append(t) |
| valid_idxs.append(i) |
|
|
| if not valid_idxs: |
| return [], "" |
|
|
| pairs = [[q_for_retrieval, cand_texts[j]] for j in range(len(valid_idxs))] |
| scores = np.array(reranker.compute_score(pairs), dtype=np.float32) |
| order = np.argsort(-scores)[:top_k] |
| |
| top = [] |
| for j in order: |
| idx = valid_idxs[j] |
| md = meta[idx].get("metadata", {}) if isinstance(meta[idx], dict) else {} |
| top.append({ |
| "idx": idx, |
| "text": cand_texts[j], |
| "rerank_score": float(scores[j]), |
| "metadata": md, |
| }) |
| |
| ctx = "\n\n---\n\n".join([t["text"] for t in top]) |
| return top, ctx |
|
|
