326 lines
11 KiB
Python
326 lines
11 KiB
Python
from __future__ import annotations
|
||
import json
|
||
import time
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import List, Tuple
|
||
|
||
import faiss
|
||
import numpy as np
|
||
import pandas as pd
|
||
|
||
# -----------------------------------------------------------
|
||
# User CONFIG – edit these values only
|
||
# -----------------------------------------------------------
|
||
|
||
# Base project path (relative to repo root)
|
||
BASE_DIR = Path(__file__).resolve().parent
|
||
|
||
# Input folder: Parquet embeddings
|
||
PATH_TO_EMBEDDED = BASE_DIR / "data" / "embeddings"
|
||
|
||
# Output folder: FAISS index + metadata
|
||
PATH_TO_FAISS = BASE_DIR / "data" / "faiss"
|
||
|
||
# Dataset version tag (example)
|
||
DATASET_VERSION = "v1.0.0"
|
||
|
||
# Index settings
|
||
INDEX_TYPE = "flat" # "flat" (exact) or "hnsw" (ANN)
|
||
METRIC = "ip" # "ip" (cosine-like, with normalization) or "l2"
|
||
HNSW_M = 32 # used only if INDEX_TYPE="hnsw"
|
||
HNSW_EFC = 200 # used only if INDEX_TYPE="hnsw"
|
||
|
||
# Meta columns to include if present in Parquet
|
||
META_COLS_PREFERRED = [
|
||
"doc_id", "text", "chapter", "section", "page",
|
||
"source_path", "hash", "model_name", "dim", "created_at",
|
||
]
|
||
|
||
# --- Single-file mode (used when RUN_MODE="single") ---
|
||
FILE_NAME = "" # without .parquet (leave empty for dir mode)
|
||
|
||
# --- Directory mode (used when RUN_MODE="dir") ---
|
||
GLOB_PATTERN = "*.parquet" # which files to index
|
||
MERGE_INTO_SINGLE = False # False=one index per file; True=one merged index
|
||
MERGED_NAME = "merged" # used only if MERGE_INTO_SINGLE=True
|
||
|
||
# Which action to run when executing this file directly:
|
||
RUN_MODE = "dir" # "single" / "dir" / "none"
|
||
# -----------------------------------------------------------
|
||
|
||
# ============ Logging setup ============
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
datefmt="%H:%M:%S"
|
||
)
|
||
|
||
# ============ Helpers ============
|
||
def _ensure_dir(path: Path) -> None:
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
|
||
def _stack_embeddings(df: pd.DataFrame) -> Tuple[np.ndarray, int]:
|
||
"""
|
||
Ensure 'embedding' column exists and stack to (N, D) float32 matrix.
|
||
Returns: (embeddings, dim)
|
||
"""
|
||
if "embedding" not in df.columns:
|
||
raise ValueError("Missing 'embedding' column in the Parquet file.")
|
||
|
||
dims = {len(v) for v in df["embedding"].head(100) if isinstance(v, (list, tuple, np.ndarray))}
|
||
if len(dims) != 1:
|
||
raise ValueError(f"Inconsistent embedding dims in sample: {dims}")
|
||
dim = list(dims)[0]
|
||
|
||
try:
|
||
emb = np.vstack(df["embedding"].values).astype("float32")
|
||
except Exception as e:
|
||
raise ValueError(f"Failed to stack embeddings to 2D float32 array: {e}")
|
||
if emb.ndim != 2 or emb.shape[1] != dim:
|
||
raise ValueError(f"Unexpected embeddings shape {emb.shape}, expected (N, {dim})")
|
||
return emb, dim
|
||
|
||
def _build_faiss_index(
|
||
emb: np.ndarray,
|
||
*,
|
||
metric: str = "ip", # "ip" (cosine-like via L2 normalization) or "l2"
|
||
index_type: str = "flat", # "flat" or "hnsw"
|
||
hnsw_m: int = 32,
|
||
hnsw_efc: int = 200,
|
||
) -> faiss.Index:
|
||
"""
|
||
Build a FAISS index for the given embeddings matrix (N, D).
|
||
- metric="ip": we L2-normalize embeddings -> inner-product ~ cosine
|
||
- index_type="flat": exact search; "hnsw": fast ANN for large N
|
||
"""
|
||
d = emb.shape[1]
|
||
if metric not in ("ip", "l2"):
|
||
raise ValueError("metric must be 'ip' or 'l2'")
|
||
|
||
if metric == "ip":
|
||
faiss.normalize_L2(emb) # cosine-like scoring
|
||
|
||
if index_type == "flat":
|
||
index = faiss.IndexFlatIP(d) if metric == "ip" else faiss.IndexFlatL2(d)
|
||
elif index_type == "hnsw":
|
||
index = faiss.IndexHNSWFlat(d, hnsw_m)
|
||
index.hnsw.efConstruction = hnsw_efc
|
||
else:
|
||
raise ValueError("index_type must be 'flat' or 'hnsw'")
|
||
|
||
index.add(emb)
|
||
return index
|
||
|
||
def _write_artifacts(
|
||
index: faiss.Index,
|
||
meta_records: List[dict],
|
||
out_index_path: Path,
|
||
out_meta_path: Path,
|
||
meta_header: dict,
|
||
) -> Tuple[Path, Path]:
|
||
faiss.write_index(index, str(out_index_path))
|
||
with out_meta_path.open("w", encoding="utf-8") as f:
|
||
json.dump({**meta_header, "records": meta_records}, f, ensure_ascii=False)
|
||
return out_index_path, out_meta_path
|
||
|
||
|
||
# ============ Public API ============
|
||
|
||
def build_index_from_parquet_for_single_file(
|
||
path_to_embedded: str,
|
||
path_to_faiss: str,
|
||
file_name: str,
|
||
*,
|
||
dataset_version: str = "unknown",
|
||
index_type: str = "flat", # "flat" or "hnsw"
|
||
metric: str = "ip", # "ip" (cosine-like) or "l2"
|
||
meta_cols_preferred: List[str] | None = None,
|
||
) -> Tuple[str, str]:
|
||
"""
|
||
Build a FAISS index from a single Parquet file that already contains an 'embedding' column.
|
||
"""
|
||
t0 = time.time()
|
||
meta_cols_preferred = meta_cols_preferred or META_COLS_PREFERRED
|
||
|
||
embedded_dir = Path(path_to_embedded)
|
||
out_dir = Path(path_to_faiss)
|
||
_ensure_dir(out_dir)
|
||
|
||
parquet_path = embedded_dir / f"{file_name}.parquet"
|
||
if not parquet_path.exists():
|
||
raise FileNotFoundError(f"Parquet not found: {parquet_path}")
|
||
|
||
logging.info(f"Loading Parquet: {parquet_path}")
|
||
df = pd.read_parquet(parquet_path)
|
||
logging.info(f"Rows: {len(df):,}")
|
||
|
||
emb, dim = _stack_embeddings(df)
|
||
logging.info(f"Embedding dim: {dim}")
|
||
|
||
logging.info(f"Building FAISS index (type={index_type}, metric={metric}) ...")
|
||
index = _build_faiss_index(
|
||
emb, metric=metric, index_type=index_type, hnsw_m=HNSW_M, hnsw_efc=HNSW_EFC
|
||
)
|
||
|
||
index_path = out_dir / f"{file_name}.index"
|
||
meta_path = out_dir / f"{file_name}_meta.json"
|
||
|
||
meta_cols = [c for c in meta_cols_preferred if c in df.columns]
|
||
meta_records = df[meta_cols].to_dict(orient="records")
|
||
|
||
header = {
|
||
"dataset_version": dataset_version,
|
||
"index_name": file_name,
|
||
"index_type": index_type,
|
||
"metric": metric,
|
||
"dim": dim,
|
||
"rows": len(df),
|
||
"meta_cols": meta_cols,
|
||
"source_parquet": str(parquet_path),
|
||
}
|
||
|
||
idx_p, meta_p = _write_artifacts(index, meta_records, index_path, meta_path, header)
|
||
dt = time.time() - t0
|
||
logging.info(f"✅ Wrote index: {idx_p}")
|
||
logging.info(f"✅ Wrote meta : {meta_p}")
|
||
logging.info(f"Done in {dt:.2f}s")
|
||
return str(idx_p), str(meta_p)
|
||
|
||
|
||
def build_index_from_parquet_for_files(
|
||
path_to_embedded: str,
|
||
path_to_faiss: str,
|
||
*,
|
||
dataset_version: str = "unknown",
|
||
index_type: str = "flat",
|
||
metric: str = "ip",
|
||
glob_pattern: str = "*.parquet",
|
||
merge_into_single: bool = False,
|
||
merged_name: str = "merged",
|
||
meta_cols_preferred: List[str] | None = None,
|
||
) -> List[Tuple[str, str]]:
|
||
"""
|
||
Build FAISS indices for all Parquet files under `path_to_embedded`.
|
||
"""
|
||
meta_cols_preferred = meta_cols_preferred or META_COLS_PREFERRED
|
||
|
||
embedded_dir = Path(path_to_embedded)
|
||
out_dir = Path(path_to_faiss)
|
||
_ensure_dir(out_dir)
|
||
|
||
parquet_paths = sorted(embedded_dir.glob(glob_pattern))
|
||
total = len(parquet_paths)
|
||
if total == 0:
|
||
raise FileNotFoundError(f"No parquet files under: {embedded_dir} (pattern={glob_pattern})")
|
||
|
||
results: List[Tuple[str, str]] = []
|
||
|
||
if merge_into_single:
|
||
logging.info(f"Merging {total} parquet files into one FAISS index: {merged_name}")
|
||
t0 = time.time()
|
||
frames = []
|
||
for i, p in enumerate(parquet_paths, start=1):
|
||
logging.info(f"[{i}/{total}] Reading: {p.name}")
|
||
df = pd.read_parquet(p)
|
||
if "embedding" not in df.columns:
|
||
raise ValueError(f"Missing 'embedding' column in {p}")
|
||
frames.append(df)
|
||
|
||
df_all = pd.concat(frames, ignore_index=True)
|
||
emb, dim = _stack_embeddings(df_all)
|
||
logging.info(f"Total rows: {len(df_all):,} | dim: {dim}")
|
||
|
||
logging.info(f"Building FAISS index (type={index_type}, metric={metric}) ...")
|
||
index = _build_faiss_index(
|
||
emb, metric=metric, index_type=index_type, hnsw_m=HNSW_M, hnsw_efc=HNSW_EFC
|
||
)
|
||
|
||
index_path = out_dir / f"{merged_name}.index"
|
||
meta_path = out_dir / f"{merged_name}_meta.json"
|
||
|
||
meta_cols = [c for c in meta_cols_preferred if c in df_all.columns]
|
||
meta_records = df_all[meta_cols].to_dict(orient="records")
|
||
header = {
|
||
"dataset_version": dataset_version,
|
||
"index_name": merged_name,
|
||
"index_type": index_type,
|
||
"metric": metric,
|
||
"dim": dim,
|
||
"rows": len(df_all),
|
||
"meta_cols": meta_cols,
|
||
"sources": [str(p) for p in parquet_paths],
|
||
}
|
||
|
||
_write_artifacts(index, meta_records, index_path, meta_path, header)
|
||
dt = time.time() - t0
|
||
logging.info(f"✅ Wrote merged index: {index_path}")
|
||
logging.info(f"✅ Wrote merged meta : {meta_path}")
|
||
logging.info(f"Done in {dt:.2f}s")
|
||
results.append((str(index_path), str(meta_path)))
|
||
return results
|
||
|
||
# One FAISS index per Parquet file, with % progress and ETA
|
||
start = time.time()
|
||
for i, p in enumerate(parquet_paths, start=1):
|
||
t_loop = time.time()
|
||
pct = (i / total) * 100.0
|
||
logging.info(f"[{i}/{total}] ({pct:5.1f}%) Processing: {p.name}")
|
||
|
||
try:
|
||
idx_p, meta_p = build_index_from_parquet_for_single_file(
|
||
path_to_embedded=str(embedded_dir),
|
||
path_to_faiss=str(out_dir),
|
||
file_name=p.stem,
|
||
dataset_version=dataset_version,
|
||
index_type=index_type,
|
||
metric=metric,
|
||
meta_cols_preferred=meta_cols_preferred,
|
||
)
|
||
results.append((idx_p, meta_p))
|
||
except Exception as e:
|
||
logging.error(f"⚠️ Skipped {p.name} due to error: {e}")
|
||
finally:
|
||
# ETA
|
||
elapsed_total = time.time() - start
|
||
avg_per_file = elapsed_total / i
|
||
remaining = total - i
|
||
eta_sec = avg_per_file * remaining
|
||
logging.info(
|
||
f"Elapsed: {elapsed_total:.1f}s | Avg/file: {avg_per_file:.2f}s | ETA: {eta_sec:.1f}s"
|
||
)
|
||
logging.info("-" * 60)
|
||
|
||
logging.info(f"✅ Completed building FAISS indexes for {len(results)} files "
|
||
f"(out of {total}). Output dir: {out_dir}")
|
||
return results
|
||
|
||
|
||
# ============ Optional direct run using CONFIG ============
|
||
|
||
if __name__ == "__main__":
|
||
if RUN_MODE == "single":
|
||
build_index_from_parquet_for_single_file(
|
||
path_to_embedded=PATH_TO_EMBEDDED,
|
||
path_to_faiss=PATH_TO_FAISS,
|
||
file_name=FILE_NAME,
|
||
dataset_version=DATASET_VERSION,
|
||
index_type=INDEX_TYPE,
|
||
metric=METRIC,
|
||
meta_cols_preferred=META_COLS_PREFERRED,
|
||
)
|
||
elif RUN_MODE == "dir":
|
||
build_index_from_parquet_for_files(
|
||
path_to_embedded=PATH_TO_EMBEDDED,
|
||
path_to_faiss=PATH_TO_FAISS,
|
||
dataset_version=DATASET_VERSION,
|
||
index_type=INDEX_TYPE,
|
||
metric=METRIC,
|
||
glob_pattern=GLOB_PATTERN,
|
||
merge_into_single=MERGE_INTO_SINGLE,
|
||
merged_name=MERGED_NAME,
|
||
meta_cols_preferred=META_COLS_PREFERRED,
|
||
)
|
||
else:
|
||
logging.info("Configured RUN_MODE='none'. Edit the CONFIG at the top and set RUN_MODE to 'single' or 'dir'.")
|