Compare commits
8 Commits
colqwen
...
dynamic-ad
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d02aee6901 | ||
|
|
62a5d7b31d | ||
|
|
0a69118f87 | ||
|
|
880a039e1d | ||
|
|
4a39b40e72 | ||
|
|
ed5fd88a85 | ||
|
|
8f4f2b4873 | ||
|
|
6a06bd893a |
5
.gitignore
vendored
5
.gitignore
vendored
@@ -18,7 +18,6 @@ demo/experiment_results/**/*.json
|
|||||||
*.eml
|
*.eml
|
||||||
*.emlx
|
*.emlx
|
||||||
*.json
|
*.json
|
||||||
*.png
|
|
||||||
!.vscode/*.json
|
!.vscode/*.json
|
||||||
*.sh
|
*.sh
|
||||||
*.txt
|
*.txt
|
||||||
@@ -102,6 +101,4 @@ CLAUDE.local.md
|
|||||||
.claude/*.local.*
|
.claude/*.local.*
|
||||||
.claude/local/*
|
.claude/local/*
|
||||||
benchmarks/data/
|
benchmarks/data/
|
||||||
|
test_add/*
|
||||||
## multi vector
|
|
||||||
apps/multimodal/vision-based-pdf-multi-vector/multi-vector-colpali-native-weaviate.py
|
|
||||||
|
|||||||
@@ -1,182 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
|
|
||||||
def _ensure_repo_paths_importable(current_file: str) -> None:
|
|
||||||
_repo_root = Path(current_file).resolve().parents[3]
|
|
||||||
_leann_core_src = _repo_root / "packages" / "leann-core" / "src"
|
|
||||||
_leann_hnsw_pkg = _repo_root / "packages" / "leann-backend-hnsw"
|
|
||||||
if str(_leann_core_src) not in sys.path:
|
|
||||||
sys.path.append(str(_leann_core_src))
|
|
||||||
if str(_leann_hnsw_pkg) not in sys.path:
|
|
||||||
sys.path.append(str(_leann_hnsw_pkg))
|
|
||||||
|
|
||||||
|
|
||||||
_ensure_repo_paths_importable(__file__)
|
|
||||||
|
|
||||||
from leann_backend_hnsw.hnsw_backend import HNSWBuilder, HNSWSearcher # noqa: E402
|
|
||||||
|
|
||||||
|
|
||||||
class LeannMultiVector:
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
index_path: str,
|
|
||||||
dim: int = 128,
|
|
||||||
distance_metric: str = "mips",
|
|
||||||
m: int = 16,
|
|
||||||
ef_construction: int = 500,
|
|
||||||
is_compact: bool = False,
|
|
||||||
is_recompute: bool = False,
|
|
||||||
embedding_model_name: str = "colvision",
|
|
||||||
) -> None:
|
|
||||||
self.index_path = index_path
|
|
||||||
self.dim = dim
|
|
||||||
self.embedding_model_name = embedding_model_name
|
|
||||||
self._pending_items: list[dict] = []
|
|
||||||
self._backend_kwargs = {
|
|
||||||
"distance_metric": distance_metric,
|
|
||||||
"M": m,
|
|
||||||
"efConstruction": ef_construction,
|
|
||||||
"is_compact": is_compact,
|
|
||||||
"is_recompute": is_recompute,
|
|
||||||
}
|
|
||||||
self._labels_meta: list[dict] = []
|
|
||||||
|
|
||||||
def _meta_dict(self) -> dict:
|
|
||||||
return {
|
|
||||||
"version": "1.0",
|
|
||||||
"backend_name": "hnsw",
|
|
||||||
"embedding_model": self.embedding_model_name,
|
|
||||||
"embedding_mode": "custom",
|
|
||||||
"dimensions": self.dim,
|
|
||||||
"backend_kwargs": self._backend_kwargs,
|
|
||||||
"is_compact": self._backend_kwargs.get("is_compact", True),
|
|
||||||
"is_pruned": self._backend_kwargs.get("is_compact", True)
|
|
||||||
and self._backend_kwargs.get("is_recompute", True),
|
|
||||||
}
|
|
||||||
|
|
||||||
def create_collection(self) -> None:
|
|
||||||
path = Path(self.index_path)
|
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
def insert(self, data: dict) -> None:
|
|
||||||
self._pending_items.append(
|
|
||||||
{
|
|
||||||
"doc_id": int(data["doc_id"]),
|
|
||||||
"filepath": data.get("filepath", ""),
|
|
||||||
"colbert_vecs": [np.asarray(v, dtype=np.float32) for v in data["colbert_vecs"]],
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
def _labels_path(self) -> Path:
|
|
||||||
index_path_obj = Path(self.index_path)
|
|
||||||
return index_path_obj.parent / f"{index_path_obj.name}.labels.json"
|
|
||||||
|
|
||||||
def _meta_path(self) -> Path:
|
|
||||||
index_path_obj = Path(self.index_path)
|
|
||||||
return index_path_obj.parent / f"{index_path_obj.name}.meta.json"
|
|
||||||
|
|
||||||
def create_index(self) -> None:
|
|
||||||
if not self._pending_items:
|
|
||||||
return
|
|
||||||
|
|
||||||
embeddings: list[np.ndarray] = []
|
|
||||||
labels_meta: list[dict] = []
|
|
||||||
|
|
||||||
for item in self._pending_items:
|
|
||||||
doc_id = int(item["doc_id"])
|
|
||||||
filepath = item.get("filepath", "")
|
|
||||||
colbert_vecs = item["colbert_vecs"]
|
|
||||||
for seq_id, vec in enumerate(colbert_vecs):
|
|
||||||
vec_np = np.asarray(vec, dtype=np.float32)
|
|
||||||
embeddings.append(vec_np)
|
|
||||||
labels_meta.append(
|
|
||||||
{
|
|
||||||
"id": f"{doc_id}:{seq_id}",
|
|
||||||
"doc_id": doc_id,
|
|
||||||
"seq_id": int(seq_id),
|
|
||||||
"filepath": filepath,
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
if not embeddings:
|
|
||||||
return
|
|
||||||
|
|
||||||
embeddings_np = np.vstack(embeddings).astype(np.float32)
|
|
||||||
# print shape of embeddings_np
|
|
||||||
print(embeddings_np.shape)
|
|
||||||
|
|
||||||
builder = HNSWBuilder(**{**self._backend_kwargs, "dimensions": self.dim})
|
|
||||||
ids = [str(i) for i in range(embeddings_np.shape[0])]
|
|
||||||
builder.build(embeddings_np, ids, self.index_path)
|
|
||||||
|
|
||||||
import json as _json
|
|
||||||
|
|
||||||
with open(self._meta_path(), "w", encoding="utf-8") as f:
|
|
||||||
_json.dump(self._meta_dict(), f, indent=2)
|
|
||||||
with open(self._labels_path(), "w", encoding="utf-8") as f:
|
|
||||||
_json.dump(labels_meta, f)
|
|
||||||
|
|
||||||
self._labels_meta = labels_meta
|
|
||||||
|
|
||||||
def _load_labels_meta_if_needed(self) -> None:
|
|
||||||
if self._labels_meta:
|
|
||||||
return
|
|
||||||
labels_path = self._labels_path()
|
|
||||||
if labels_path.exists():
|
|
||||||
import json as _json
|
|
||||||
|
|
||||||
with open(labels_path, encoding="utf-8") as f:
|
|
||||||
self._labels_meta = _json.load(f)
|
|
||||||
|
|
||||||
def search(
|
|
||||||
self, data: np.ndarray, topk: int, first_stage_k: int = 50
|
|
||||||
) -> list[tuple[float, int]]:
|
|
||||||
if data.ndim == 1:
|
|
||||||
data = data.reshape(1, -1)
|
|
||||||
if data.dtype != np.float32:
|
|
||||||
data = data.astype(np.float32)
|
|
||||||
|
|
||||||
self._load_labels_meta_if_needed()
|
|
||||||
|
|
||||||
searcher = HNSWSearcher(self.index_path, meta=self._meta_dict())
|
|
||||||
raw = searcher.search(
|
|
||||||
data,
|
|
||||||
first_stage_k,
|
|
||||||
recompute_embeddings=False,
|
|
||||||
complexity=128,
|
|
||||||
beam_width=1,
|
|
||||||
prune_ratio=0.0,
|
|
||||||
batch_size=0,
|
|
||||||
)
|
|
||||||
|
|
||||||
labels = raw.get("labels")
|
|
||||||
distances = raw.get("distances")
|
|
||||||
if labels is None or distances is None:
|
|
||||||
return []
|
|
||||||
|
|
||||||
doc_scores: dict[int, float] = {}
|
|
||||||
B = len(labels)
|
|
||||||
for b in range(B):
|
|
||||||
per_doc_best: dict[int, float] = {}
|
|
||||||
for k, sid in enumerate(labels[b]):
|
|
||||||
try:
|
|
||||||
idx = int(sid)
|
|
||||||
except Exception:
|
|
||||||
continue
|
|
||||||
if 0 <= idx < len(self._labels_meta):
|
|
||||||
doc_id = int(self._labels_meta[idx]["doc_id"]) # type: ignore[index]
|
|
||||||
else:
|
|
||||||
continue
|
|
||||||
score = float(distances[b][k])
|
|
||||||
if (doc_id not in per_doc_best) or (score > per_doc_best[doc_id]):
|
|
||||||
per_doc_best[doc_id] = score
|
|
||||||
for doc_id, best_score in per_doc_best.items():
|
|
||||||
doc_scores[doc_id] = doc_scores.get(doc_id, 0.0) + best_score
|
|
||||||
|
|
||||||
scores = sorted(((v, k) for k, v in doc_scores.items()), key=lambda x: x[0], reverse=True)
|
|
||||||
return scores[:topk] if len(scores) >= topk else scores
|
|
||||||
@@ -1,477 +0,0 @@
|
|||||||
## Jupyter-style notebook script
|
|
||||||
# %%
|
|
||||||
# uv pip install matplotlib qwen_vl_utils
|
|
||||||
import os
|
|
||||||
import re
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any, Optional, cast
|
|
||||||
|
|
||||||
from PIL import Image
|
|
||||||
from tqdm import tqdm
|
|
||||||
|
|
||||||
|
|
||||||
def _ensure_repo_paths_importable(current_file: str) -> None:
|
|
||||||
"""Make local leann packages importable without installing (mirrors multi-vector-leann.py)."""
|
|
||||||
_repo_root = Path(current_file).resolve().parents[3]
|
|
||||||
_leann_core_src = _repo_root / "packages" / "leann-core" / "src"
|
|
||||||
_leann_hnsw_pkg = _repo_root / "packages" / "leann-backend-hnsw"
|
|
||||||
if str(_leann_core_src) not in sys.path:
|
|
||||||
sys.path.append(str(_leann_core_src))
|
|
||||||
if str(_leann_hnsw_pkg) not in sys.path:
|
|
||||||
sys.path.append(str(_leann_hnsw_pkg))
|
|
||||||
|
|
||||||
|
|
||||||
_ensure_repo_paths_importable(__file__)
|
|
||||||
|
|
||||||
from leann_multi_vector import LeannMultiVector # noqa: E402
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Config
|
|
||||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
||||||
QUERY = "How does DeepSeek-V2 compare against the LLaMA family of LLMs?"
|
|
||||||
MODEL: str = "colqwen2" # "colpali" or "colqwen2"
|
|
||||||
|
|
||||||
# Data source: set to True to use the Hugging Face dataset example (recommended)
|
|
||||||
USE_HF_DATASET: bool = True
|
|
||||||
DATASET_NAME: str = "weaviate/arXiv-AI-papers-multi-vector"
|
|
||||||
DATASET_SPLIT: str = "train"
|
|
||||||
MAX_DOCS: Optional[int] = None # limit number of pages to index; None = all
|
|
||||||
|
|
||||||
# Local pages (used when USE_HF_DATASET == False)
|
|
||||||
PDF: Optional[str] = None # e.g., "./pdfs/2004.12832v2.pdf"
|
|
||||||
PAGES_DIR: str = "./pages"
|
|
||||||
|
|
||||||
# Index + retrieval settings
|
|
||||||
INDEX_PATH: str = "./indexes/colvision.leann"
|
|
||||||
TOPK: int = 1
|
|
||||||
FIRST_STAGE_K: int = 500
|
|
||||||
REBUILD_INDEX: bool = False
|
|
||||||
|
|
||||||
# Artifacts
|
|
||||||
SAVE_TOP_IMAGE: Optional[str] = "./figures/retrieved_page.png"
|
|
||||||
SIMILARITY_MAP: bool = True
|
|
||||||
SIM_TOKEN_IDX: int = 13 # -1 means auto-select the most salient token
|
|
||||||
SIM_OUTPUT: str = "./figures/similarity_map.png"
|
|
||||||
ANSWER: bool = True
|
|
||||||
MAX_NEW_TOKENS: int = 128
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Helpers
|
|
||||||
def _natural_sort_key(name: str) -> int:
|
|
||||||
m = re.search(r"\d+", name)
|
|
||||||
return int(m.group()) if m else 0
|
|
||||||
|
|
||||||
|
|
||||||
def _load_images_from_dir(pages_dir: str) -> tuple[list[str], list[Image.Image]]:
|
|
||||||
filenames = [n for n in os.listdir(pages_dir) if n.lower().endswith((".png", ".jpg", ".jpeg"))]
|
|
||||||
filenames = sorted(filenames, key=_natural_sort_key)
|
|
||||||
filepaths = [os.path.join(pages_dir, n) for n in filenames]
|
|
||||||
images = [Image.open(p) for p in filepaths]
|
|
||||||
return filepaths, images
|
|
||||||
|
|
||||||
|
|
||||||
def _maybe_convert_pdf_to_images(pdf_path: Optional[str], pages_dir: str, dpi: int = 200) -> None:
|
|
||||||
if not pdf_path:
|
|
||||||
return
|
|
||||||
os.makedirs(pages_dir, exist_ok=True)
|
|
||||||
try:
|
|
||||||
from pdf2image import convert_from_path
|
|
||||||
except Exception as e:
|
|
||||||
raise RuntimeError(
|
|
||||||
"pdf2image is required to convert PDF to images. Install via pip install pdf2image"
|
|
||||||
) from e
|
|
||||||
images = convert_from_path(pdf_path, dpi=dpi)
|
|
||||||
for i, image in enumerate(images):
|
|
||||||
image.save(os.path.join(pages_dir, f"page_{i + 1}.png"), "PNG")
|
|
||||||
|
|
||||||
|
|
||||||
def _select_device_and_dtype():
|
|
||||||
import torch
|
|
||||||
from colpali_engine.utils.torch_utils import get_torch_device
|
|
||||||
|
|
||||||
device_str = (
|
|
||||||
"cuda"
|
|
||||||
if torch.cuda.is_available()
|
|
||||||
else (
|
|
||||||
"mps"
|
|
||||||
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available()
|
|
||||||
else "cpu"
|
|
||||||
)
|
|
||||||
)
|
|
||||||
device = get_torch_device(device_str)
|
|
||||||
# Stable dtype selection to avoid NaNs:
|
|
||||||
# - CUDA: prefer bfloat16 if supported, else float16
|
|
||||||
# - MPS: use float32 (fp16 on MPS can produce NaNs in some ops)
|
|
||||||
# - CPU: float32
|
|
||||||
if device_str == "cuda":
|
|
||||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
|
||||||
try:
|
|
||||||
torch.backends.cuda.matmul.allow_tf32 = True # Better stability/perf on Ampere+
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
elif device_str == "mps":
|
|
||||||
dtype = torch.float32
|
|
||||||
else:
|
|
||||||
dtype = torch.float32
|
|
||||||
return device_str, device, dtype
|
|
||||||
|
|
||||||
|
|
||||||
def _load_colvision(model_choice: str):
|
|
||||||
import torch
|
|
||||||
from colpali_engine.models import ColPali, ColQwen2, ColQwen2Processor
|
|
||||||
from colpali_engine.models.paligemma.colpali.processing_colpali import ColPaliProcessor
|
|
||||||
from transformers.utils.import_utils import is_flash_attn_2_available
|
|
||||||
|
|
||||||
device_str, device, dtype = _select_device_and_dtype()
|
|
||||||
|
|
||||||
if model_choice == "colqwen2":
|
|
||||||
model_name = "vidore/colqwen2-v1.0"
|
|
||||||
# On CPU/MPS we must avoid flash-attn and stay eager; on CUDA prefer flash-attn if available
|
|
||||||
attn_implementation = (
|
|
||||||
"flash_attention_2"
|
|
||||||
if (device_str == "cuda" and is_flash_attn_2_available())
|
|
||||||
else "eager"
|
|
||||||
)
|
|
||||||
model = ColQwen2.from_pretrained(
|
|
||||||
model_name,
|
|
||||||
torch_dtype=torch.bfloat16,
|
|
||||||
device_map=device,
|
|
||||||
attn_implementation=attn_implementation,
|
|
||||||
).eval()
|
|
||||||
processor = ColQwen2Processor.from_pretrained(model_name)
|
|
||||||
else:
|
|
||||||
model_name = "vidore/colpali-v1.2"
|
|
||||||
model = ColPali.from_pretrained(
|
|
||||||
model_name,
|
|
||||||
torch_dtype=torch.bfloat16,
|
|
||||||
device_map=device,
|
|
||||||
).eval()
|
|
||||||
processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
|
|
||||||
|
|
||||||
return model_name, model, processor, device_str, device, dtype
|
|
||||||
|
|
||||||
|
|
||||||
def _embed_images(model, processor, images: list[Image.Image]) -> list[Any]:
|
|
||||||
import torch
|
|
||||||
from colpali_engine.utils.torch_utils import ListDataset
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
|
|
||||||
# Ensure deterministic eval and autocast for stability
|
|
||||||
model.eval()
|
|
||||||
|
|
||||||
dataloader = DataLoader(
|
|
||||||
dataset=ListDataset[Image.Image](images),
|
|
||||||
batch_size=1,
|
|
||||||
shuffle=False,
|
|
||||||
collate_fn=lambda x: processor.process_images(x),
|
|
||||||
)
|
|
||||||
|
|
||||||
doc_vecs: list[Any] = []
|
|
||||||
for batch_doc in dataloader:
|
|
||||||
with torch.no_grad():
|
|
||||||
batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
|
|
||||||
# autocast on CUDA for bf16/fp16; on CPU/MPS stay in fp32
|
|
||||||
if model.device.type == "cuda":
|
|
||||||
with torch.autocast(
|
|
||||||
device_type="cuda",
|
|
||||||
dtype=model.dtype if model.dtype.is_floating_point else torch.bfloat16,
|
|
||||||
):
|
|
||||||
embeddings_doc = model(**batch_doc)
|
|
||||||
else:
|
|
||||||
embeddings_doc = model(**batch_doc)
|
|
||||||
doc_vecs.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
|
|
||||||
return doc_vecs
|
|
||||||
|
|
||||||
|
|
||||||
def _embed_queries(model, processor, queries: list[str]) -> list[Any]:
|
|
||||||
import torch
|
|
||||||
from colpali_engine.utils.torch_utils import ListDataset
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
|
|
||||||
model.eval()
|
|
||||||
|
|
||||||
dataloader = DataLoader(
|
|
||||||
dataset=ListDataset[str](queries),
|
|
||||||
batch_size=1,
|
|
||||||
shuffle=False,
|
|
||||||
collate_fn=lambda x: processor.process_queries(x),
|
|
||||||
)
|
|
||||||
|
|
||||||
q_vecs: list[Any] = []
|
|
||||||
for batch_query in dataloader:
|
|
||||||
with torch.no_grad():
|
|
||||||
batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
|
|
||||||
if model.device.type == "cuda":
|
|
||||||
with torch.autocast(
|
|
||||||
device_type="cuda",
|
|
||||||
dtype=model.dtype if model.dtype.is_floating_point else torch.bfloat16,
|
|
||||||
):
|
|
||||||
embeddings_query = model(**batch_query)
|
|
||||||
else:
|
|
||||||
embeddings_query = model(**batch_query)
|
|
||||||
q_vecs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
|
|
||||||
return q_vecs
|
|
||||||
|
|
||||||
|
|
||||||
def _build_index(index_path: str, doc_vecs: list[Any], filepaths: list[str]) -> LeannMultiVector:
|
|
||||||
dim = int(doc_vecs[0].shape[-1])
|
|
||||||
retriever = LeannMultiVector(index_path=index_path, dim=dim)
|
|
||||||
retriever.create_collection()
|
|
||||||
for i, vec in enumerate(doc_vecs):
|
|
||||||
data = {
|
|
||||||
"colbert_vecs": vec.float().numpy(),
|
|
||||||
"doc_id": i,
|
|
||||||
"filepath": filepaths[i],
|
|
||||||
}
|
|
||||||
retriever.insert(data)
|
|
||||||
retriever.create_index()
|
|
||||||
return retriever
|
|
||||||
|
|
||||||
|
|
||||||
def _load_retriever_if_index_exists(index_path: str, dim: int) -> Optional[LeannMultiVector]:
|
|
||||||
index_base = Path(index_path)
|
|
||||||
# Rough heuristic: index dir exists AND meta+labels files exist
|
|
||||||
meta = index_base.parent / f"{index_base.name}.meta.json"
|
|
||||||
labels = index_base.parent / f"{index_base.name}.labels.json"
|
|
||||||
if index_base.exists() and meta.exists() and labels.exists():
|
|
||||||
return LeannMultiVector(index_path=index_path, dim=dim)
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def _generate_similarity_map(
|
|
||||||
model,
|
|
||||||
processor,
|
|
||||||
image: Image.Image,
|
|
||||||
query: str,
|
|
||||||
token_idx: Optional[int] = None,
|
|
||||||
output_path: Optional[str] = None,
|
|
||||||
) -> tuple[int, float]:
|
|
||||||
import torch
|
|
||||||
from colpali_engine.interpretability import (
|
|
||||||
get_similarity_maps_from_embeddings,
|
|
||||||
plot_similarity_map,
|
|
||||||
)
|
|
||||||
|
|
||||||
batch_images = processor.process_images([image]).to(model.device)
|
|
||||||
batch_queries = processor.process_queries([query]).to(model.device)
|
|
||||||
|
|
||||||
with torch.no_grad():
|
|
||||||
image_embeddings = model.forward(**batch_images)
|
|
||||||
query_embeddings = model.forward(**batch_queries)
|
|
||||||
|
|
||||||
n_patches = processor.get_n_patches(
|
|
||||||
image_size=image.size,
|
|
||||||
spatial_merge_size=getattr(model, "spatial_merge_size", None),
|
|
||||||
)
|
|
||||||
image_mask = processor.get_image_mask(batch_images)
|
|
||||||
|
|
||||||
batched_similarity_maps = get_similarity_maps_from_embeddings(
|
|
||||||
image_embeddings=image_embeddings,
|
|
||||||
query_embeddings=query_embeddings,
|
|
||||||
n_patches=n_patches,
|
|
||||||
image_mask=image_mask,
|
|
||||||
)
|
|
||||||
|
|
||||||
similarity_maps = batched_similarity_maps[0]
|
|
||||||
|
|
||||||
# Determine token index if not provided: choose the token with highest max score
|
|
||||||
if token_idx is None:
|
|
||||||
per_token_max = similarity_maps.view(similarity_maps.shape[0], -1).max(dim=1).values
|
|
||||||
token_idx = int(per_token_max.argmax().item())
|
|
||||||
|
|
||||||
max_sim_score = similarity_maps[token_idx, :, :].max().item()
|
|
||||||
|
|
||||||
if output_path:
|
|
||||||
import matplotlib.pyplot as plt
|
|
||||||
|
|
||||||
fig, ax = plot_similarity_map(
|
|
||||||
image=image,
|
|
||||||
similarity_map=similarity_maps[token_idx],
|
|
||||||
figsize=(14, 14),
|
|
||||||
show_colorbar=False,
|
|
||||||
)
|
|
||||||
ax.set_title(f"Token #{token_idx}. MaxSim score: {max_sim_score:.2f}", fontsize=12)
|
|
||||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
||||||
plt.savefig(output_path, bbox_inches="tight")
|
|
||||||
plt.close(fig)
|
|
||||||
|
|
||||||
return token_idx, float(max_sim_score)
|
|
||||||
|
|
||||||
|
|
||||||
class QwenVL:
|
|
||||||
def __init__(self, device: str):
|
|
||||||
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
|
|
||||||
from transformers.utils.import_utils import is_flash_attn_2_available
|
|
||||||
|
|
||||||
attn_implementation = "flash_attention_2" if is_flash_attn_2_available() else "eager"
|
|
||||||
self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
||||||
"Qwen/Qwen2.5-VL-3B-Instruct",
|
|
||||||
torch_dtype="auto",
|
|
||||||
device_map=device,
|
|
||||||
attn_implementation=attn_implementation,
|
|
||||||
)
|
|
||||||
|
|
||||||
min_pixels = 256 * 28 * 28
|
|
||||||
max_pixels = 1280 * 28 * 28
|
|
||||||
self.processor = AutoProcessor.from_pretrained(
|
|
||||||
"Qwen/Qwen2.5-VL-3B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels
|
|
||||||
)
|
|
||||||
|
|
||||||
def answer(self, query: str, images: list[Image.Image], max_new_tokens: int = 128) -> str:
|
|
||||||
import base64
|
|
||||||
from io import BytesIO
|
|
||||||
|
|
||||||
from qwen_vl_utils import process_vision_info
|
|
||||||
|
|
||||||
content = []
|
|
||||||
for img in images:
|
|
||||||
buffer = BytesIO()
|
|
||||||
img.save(buffer, format="jpeg")
|
|
||||||
img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
|
||||||
content.append({"type": "image", "image": f"data:image;base64,{img_base64}"})
|
|
||||||
content.append({"type": "text", "text": query})
|
|
||||||
messages = [{"role": "user", "content": content}]
|
|
||||||
|
|
||||||
text = self.processor.apply_chat_template(
|
|
||||||
messages, tokenize=False, add_generation_prompt=True
|
|
||||||
)
|
|
||||||
image_inputs, video_inputs = process_vision_info(messages)
|
|
||||||
inputs = self.processor(
|
|
||||||
text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt"
|
|
||||||
)
|
|
||||||
inputs = inputs.to(self.model.device)
|
|
||||||
|
|
||||||
generated_ids = self.model.generate(**inputs, max_new_tokens=max_new_tokens)
|
|
||||||
generated_ids_trimmed = [
|
|
||||||
out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
|
|
||||||
]
|
|
||||||
return self.processor.batch_decode(
|
|
||||||
generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
|
|
||||||
)[0]
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
|
|
||||||
# Step 1: Prepare data
|
|
||||||
if USE_HF_DATASET:
|
|
||||||
from datasets import load_dataset
|
|
||||||
|
|
||||||
dataset = load_dataset(DATASET_NAME, split=DATASET_SPLIT)
|
|
||||||
N = len(dataset) if MAX_DOCS is None else min(MAX_DOCS, len(dataset))
|
|
||||||
filepaths: list[str] = []
|
|
||||||
images: list[Image.Image] = []
|
|
||||||
for i in tqdm(range(N), desc="Loading dataset"):
|
|
||||||
p = dataset[i]
|
|
||||||
# Compose a descriptive identifier for printing later
|
|
||||||
identifier = f"arXiv:{p['paper_arxiv_id']}|title:{p['paper_title']}|page:{int(p['page_number'])}|id:{p['page_id']}"
|
|
||||||
print(identifier)
|
|
||||||
filepaths.append(identifier)
|
|
||||||
images.append(p["page_image"]) # PIL Image
|
|
||||||
else:
|
|
||||||
_maybe_convert_pdf_to_images(PDF, PAGES_DIR)
|
|
||||||
filepaths, images = _load_images_from_dir(PAGES_DIR)
|
|
||||||
if not images:
|
|
||||||
raise RuntimeError(
|
|
||||||
f"No images found in {PAGES_DIR}. Provide PDF path in PDF variable or ensure images exist."
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Step 2: Load model and processor
|
|
||||||
model_name, model, processor, device_str, device, dtype = _load_colvision(MODEL)
|
|
||||||
print(f"Using model={model_name}, device={device_str}, dtype={dtype}")
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Step 3: Build or load index
|
|
||||||
retriever: Optional[LeannMultiVector] = None
|
|
||||||
if not REBUILD_INDEX:
|
|
||||||
try:
|
|
||||||
one_vec = _embed_images(model, processor, [images[0]])[0]
|
|
||||||
retriever = _load_retriever_if_index_exists(INDEX_PATH, dim=int(one_vec.shape[-1]))
|
|
||||||
except Exception:
|
|
||||||
retriever = None
|
|
||||||
|
|
||||||
if retriever is None:
|
|
||||||
doc_vecs = _embed_images(model, processor, images)
|
|
||||||
retriever = _build_index(INDEX_PATH, doc_vecs, filepaths)
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Step 4: Embed query and search
|
|
||||||
q_vec = _embed_queries(model, processor, [QUERY])[0]
|
|
||||||
results = retriever.search(q_vec.float().numpy(), topk=TOPK, first_stage_k=FIRST_STAGE_K)
|
|
||||||
if not results:
|
|
||||||
print("No results found.")
|
|
||||||
else:
|
|
||||||
print(f'Top {len(results)} results for query: "{QUERY}"')
|
|
||||||
top_images: list[Image.Image] = []
|
|
||||||
for rank, (score, doc_id) in enumerate(results, start=1):
|
|
||||||
path = filepaths[doc_id]
|
|
||||||
# For HF dataset, path is a descriptive identifier, not a real file path
|
|
||||||
print(f"{rank}) MaxSim: {score:.4f}, Page: {path}")
|
|
||||||
top_images.append(images[doc_id])
|
|
||||||
|
|
||||||
if SAVE_TOP_IMAGE:
|
|
||||||
from pathlib import Path as _Path
|
|
||||||
|
|
||||||
base = _Path(SAVE_TOP_IMAGE)
|
|
||||||
base.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
for rank, img in enumerate(top_images[:TOPK], start=1):
|
|
||||||
if base.suffix:
|
|
||||||
out_path = base.parent / f"{base.stem}_rank{rank}{base.suffix}"
|
|
||||||
else:
|
|
||||||
out_path = base / f"retrieved_page_rank{rank}.png"
|
|
||||||
img.save(str(out_path))
|
|
||||||
print(f"Saved retrieved page (rank {rank}) to: {out_path}")
|
|
||||||
|
|
||||||
## TODO stange results of second page of DeepSeek-V2 rather than the first page
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Step 5: Similarity maps for top-K results
|
|
||||||
if results and SIMILARITY_MAP:
|
|
||||||
token_idx = None if SIM_TOKEN_IDX < 0 else int(SIM_TOKEN_IDX)
|
|
||||||
from pathlib import Path as _Path
|
|
||||||
|
|
||||||
output_base = _Path(SIM_OUTPUT) if SIM_OUTPUT else None
|
|
||||||
for rank, img in enumerate(top_images[:TOPK], start=1):
|
|
||||||
if output_base:
|
|
||||||
if output_base.suffix:
|
|
||||||
out_dir = output_base.parent
|
|
||||||
out_name = f"{output_base.stem}_rank{rank}{output_base.suffix}"
|
|
||||||
out_path = str(out_dir / out_name)
|
|
||||||
else:
|
|
||||||
out_dir = output_base
|
|
||||||
out_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
out_path = str(out_dir / f"similarity_map_rank{rank}.png")
|
|
||||||
else:
|
|
||||||
out_path = None
|
|
||||||
chosen_idx, max_sim = _generate_similarity_map(
|
|
||||||
model=model,
|
|
||||||
processor=processor,
|
|
||||||
image=img,
|
|
||||||
query=QUERY,
|
|
||||||
token_idx=token_idx,
|
|
||||||
output_path=out_path,
|
|
||||||
)
|
|
||||||
if out_path:
|
|
||||||
print(
|
|
||||||
f"Saved similarity map for rank {rank}, token #{chosen_idx} (max={max_sim:.2f}) to: {out_path}"
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
print(
|
|
||||||
f"Computed similarity map for rank {rank}, token #{chosen_idx} (max={max_sim:.2f})"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Step 6: Optional answer generation
|
|
||||||
if results and ANSWER:
|
|
||||||
qwen = QwenVL(device=device_str)
|
|
||||||
response = qwen.answer(QUERY, top_images[:TOPK], max_new_tokens=MAX_NEW_TOKENS)
|
|
||||||
print("\nAnswer:")
|
|
||||||
print(response)
|
|
||||||
@@ -1,134 +0,0 @@
|
|||||||
# pip install pdf2image
|
|
||||||
# pip install pymilvus
|
|
||||||
# pip install colpali_engine
|
|
||||||
# pip install tqdm
|
|
||||||
# pip install pillow
|
|
||||||
|
|
||||||
# %%
|
|
||||||
from pdf2image import convert_from_path
|
|
||||||
|
|
||||||
pdf_path = "pdfs/2004.12832v2.pdf"
|
|
||||||
images = convert_from_path(pdf_path)
|
|
||||||
|
|
||||||
for i, image in enumerate(images):
|
|
||||||
image.save(f"pages/page_{i + 1}.png", "PNG")
|
|
||||||
|
|
||||||
# %%
|
|
||||||
import os
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
# Make local leann packages importable without installing
|
|
||||||
_repo_root = Path(__file__).resolve().parents[3]
|
|
||||||
_leann_core_src = _repo_root / "packages" / "leann-core" / "src"
|
|
||||||
_leann_hnsw_pkg = _repo_root / "packages" / "leann-backend-hnsw"
|
|
||||||
import sys
|
|
||||||
|
|
||||||
if str(_leann_core_src) not in sys.path:
|
|
||||||
sys.path.append(str(_leann_core_src))
|
|
||||||
if str(_leann_hnsw_pkg) not in sys.path:
|
|
||||||
sys.path.append(str(_leann_hnsw_pkg))
|
|
||||||
|
|
||||||
from leann_multi_vector import LeannMultiVector
|
|
||||||
|
|
||||||
|
|
||||||
class LeannRetriever(LeannMultiVector):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
# %%
|
|
||||||
from typing import cast
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from colpali_engine.models import ColPali
|
|
||||||
from colpali_engine.models.paligemma.colpali.processing_colpali import ColPaliProcessor
|
|
||||||
from colpali_engine.utils.torch_utils import ListDataset, get_torch_device
|
|
||||||
from torch.utils.data import DataLoader
|
|
||||||
|
|
||||||
# Auto-select device: CUDA > MPS (mac) > CPU
|
|
||||||
_device_str = (
|
|
||||||
"cuda"
|
|
||||||
if torch.cuda.is_available()
|
|
||||||
else (
|
|
||||||
"mps"
|
|
||||||
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available()
|
|
||||||
else "cpu"
|
|
||||||
)
|
|
||||||
)
|
|
||||||
device = get_torch_device(_device_str)
|
|
||||||
# Prefer fp16 on GPU/MPS, bfloat16 on CPU
|
|
||||||
_dtype = torch.float16 if _device_str in ("cuda", "mps") else torch.bfloat16
|
|
||||||
model_name = "vidore/colpali-v1.2"
|
|
||||||
|
|
||||||
model = ColPali.from_pretrained(
|
|
||||||
model_name,
|
|
||||||
torch_dtype=_dtype,
|
|
||||||
device_map=device,
|
|
||||||
).eval()
|
|
||||||
print(f"Using device={_device_str}, dtype={_dtype}")
|
|
||||||
|
|
||||||
queries = [
|
|
||||||
"How to end-to-end retrieval with ColBert",
|
|
||||||
"Where is ColBERT performance Table, including text representation results?",
|
|
||||||
]
|
|
||||||
|
|
||||||
processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained(model_name))
|
|
||||||
|
|
||||||
dataloader = DataLoader(
|
|
||||||
dataset=ListDataset[str](queries),
|
|
||||||
batch_size=1,
|
|
||||||
shuffle=False,
|
|
||||||
collate_fn=lambda x: processor.process_queries(x),
|
|
||||||
)
|
|
||||||
|
|
||||||
qs: list[torch.Tensor] = []
|
|
||||||
for batch_query in dataloader:
|
|
||||||
with torch.no_grad():
|
|
||||||
batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
|
|
||||||
embeddings_query = model(**batch_query)
|
|
||||||
qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
|
|
||||||
print(qs[0].shape)
|
|
||||||
# %%
|
|
||||||
|
|
||||||
|
|
||||||
import re
|
|
||||||
|
|
||||||
from PIL import Image
|
|
||||||
from tqdm import tqdm
|
|
||||||
|
|
||||||
page_filenames = sorted(os.listdir("./pages"), key=lambda n: int(re.search(r"\d+", n).group()))
|
|
||||||
images = [Image.open(os.path.join("./pages", name)) for name in page_filenames]
|
|
||||||
|
|
||||||
dataloader = DataLoader(
|
|
||||||
dataset=ListDataset[str](images),
|
|
||||||
batch_size=1,
|
|
||||||
shuffle=False,
|
|
||||||
collate_fn=lambda x: processor.process_images(x),
|
|
||||||
)
|
|
||||||
|
|
||||||
ds: list[torch.Tensor] = []
|
|
||||||
for batch_doc in tqdm(dataloader):
|
|
||||||
with torch.no_grad():
|
|
||||||
batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
|
|
||||||
embeddings_doc = model(**batch_doc)
|
|
||||||
ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
|
|
||||||
|
|
||||||
print(ds[0].shape)
|
|
||||||
|
|
||||||
# %%
|
|
||||||
# Build HNSW index via LeannRetriever primitives and run search
|
|
||||||
index_path = "./indexes/colpali.leann"
|
|
||||||
retriever = LeannRetriever(index_path=index_path, dim=int(ds[0].shape[-1]))
|
|
||||||
retriever.create_collection()
|
|
||||||
filepaths = [os.path.join("./pages", name) for name in page_filenames]
|
|
||||||
for i in range(len(filepaths)):
|
|
||||||
data = {
|
|
||||||
"colbert_vecs": ds[i].float().numpy(),
|
|
||||||
"doc_id": i,
|
|
||||||
"filepath": filepaths[i],
|
|
||||||
}
|
|
||||||
retriever.insert(data)
|
|
||||||
retriever.create_index()
|
|
||||||
for query in qs:
|
|
||||||
query_np = query.float().numpy()
|
|
||||||
result = retriever.search(query_np, topk=1)
|
|
||||||
print(filepaths[result[0][1]])
|
|
||||||
380
examples/dynamic_add_leann_no_recompute.py
Normal file
380
examples/dynamic_add_leann_no_recompute.py
Normal file
@@ -0,0 +1,380 @@
|
|||||||
|
"""
|
||||||
|
Dynamic add example for LEANN using HNSW backend without recompute.
|
||||||
|
|
||||||
|
- Builds a base index from a directory of documents
|
||||||
|
- Incrementally adds new documents without recomputing stored embeddings
|
||||||
|
|
||||||
|
Defaults:
|
||||||
|
- Base data: /Users/yichuan/Desktop/code/LEANN/leann/data
|
||||||
|
- Incremental data: /Users/yichuan/Desktop/code/LEANN/leann/test_add
|
||||||
|
- Index path: <index_dir>/documents.leann
|
||||||
|
|
||||||
|
Usage examples:
|
||||||
|
uv run python examples/dynamic_add_leann_no_recompute.py --build-base \
|
||||||
|
--base-dir /Users/yichuan/Desktop/code/LEANN/leann/data \
|
||||||
|
--index-dir ./test_doc_files
|
||||||
|
|
||||||
|
uv run python examples/dynamic_add_leann_no_recompute.py --add-incremental \
|
||||||
|
--add-dir /Users/yichuan/Desktop/code/LEANN/leann/test_add \
|
||||||
|
--index-dir ./test_doc_files
|
||||||
|
|
||||||
|
Quick recompute test (both true):
|
||||||
|
# Recompute build
|
||||||
|
uv run python examples/dynamic_add_leann_no_recompute.py --build-base \
|
||||||
|
--recompute-build --ef-construction 200 \
|
||||||
|
--base-dir /Users/yichuan/Desktop/code/LEANN/leann/data \
|
||||||
|
--index-dir ./test_doc_files --index-name documents.leann
|
||||||
|
|
||||||
|
# Recompute add
|
||||||
|
uv run python examples/dynamic_add_leann_no_recompute.py --add-incremental \
|
||||||
|
--recompute-add --ef-construction 32 \
|
||||||
|
--add-dir /Users/yichuan/Desktop/code/LEANN/leann/test_add \
|
||||||
|
--index-dir ./test_doc_files --index-name documents.leann
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import json
|
||||||
|
import pickle
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Optional
|
||||||
|
|
||||||
|
# Ensure we can import from the local packages and apps folders
|
||||||
|
ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
CORE_SRC = ROOT / "packages" / "leann-core" / "src"
|
||||||
|
HNSW_PKG_DIR = ROOT / "packages" / "leann-backend-hnsw"
|
||||||
|
APPS_DIR = ROOT / "apps"
|
||||||
|
|
||||||
|
|
||||||
|
# Prefer the installed backend if available (it contains the compiled extension)
|
||||||
|
def _prefer_installed(pkg_name: str) -> bool:
|
||||||
|
try:
|
||||||
|
import importlib
|
||||||
|
import importlib.util
|
||||||
|
|
||||||
|
spec = importlib.util.find_spec(pkg_name)
|
||||||
|
if spec and spec.origin and "site-packages" in spec.origin:
|
||||||
|
# ensure the faiss shim/extension is importable from the installed package
|
||||||
|
importlib.import_module(f"{pkg_name}.faiss")
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
# Prepend paths, but only add the repo backend if the installed one is not present
|
||||||
|
paths_to_prepend = [CORE_SRC, APPS_DIR]
|
||||||
|
if not _prefer_installed("leann_backend_hnsw"):
|
||||||
|
paths_to_prepend.insert(1, HNSW_PKG_DIR)
|
||||||
|
|
||||||
|
for p in paths_to_prepend:
|
||||||
|
p_str = str(p)
|
||||||
|
if p_str not in sys.path:
|
||||||
|
sys.path.insert(0, p_str)
|
||||||
|
|
||||||
|
# Defer non-stdlib imports until after sys.path setup within functions (avoid E402)
|
||||||
|
|
||||||
|
|
||||||
|
def _load_documents(data_dir: str, required_exts: Optional[list[str]] = None) -> list[Any]:
|
||||||
|
from llama_index.core import SimpleDirectoryReader # type: ignore
|
||||||
|
|
||||||
|
reader_kwargs: dict[str, Any] = {"recursive": True, "encoding": "utf-8"}
|
||||||
|
if required_exts:
|
||||||
|
reader_kwargs["required_exts"] = required_exts
|
||||||
|
documents = SimpleDirectoryReader(data_dir, **reader_kwargs).load_data(show_progress=True)
|
||||||
|
return documents
|
||||||
|
|
||||||
|
|
||||||
|
def _ensure_index_dir(index_dir: Path) -> None:
|
||||||
|
index_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
|
||||||
|
def _index_files(index_path: Path) -> tuple[Path, Path, Path]:
|
||||||
|
"""Return (passages.jsonl, passages.idx, index.index) paths for a given index base path.
|
||||||
|
|
||||||
|
Note: HNSWBackend writes the FAISS index using the stem (without .leann),
|
||||||
|
i.e., for base 'documents.leann' the file is 'documents.index'. We prefer the
|
||||||
|
existing file among candidates.
|
||||||
|
"""
|
||||||
|
passages_file = index_path.parent / f"{index_path.name}.passages.jsonl"
|
||||||
|
offsets_file = index_path.parent / f"{index_path.name}.passages.idx"
|
||||||
|
candidate_name_index = index_path.parent / f"{index_path.name}.index"
|
||||||
|
candidate_stem_index = index_path.parent / f"{index_path.stem}.index"
|
||||||
|
index_file = candidate_stem_index if candidate_stem_index.exists() else candidate_name_index
|
||||||
|
return passages_file, offsets_file, index_file
|
||||||
|
|
||||||
|
|
||||||
|
def _read_meta(index_path: Path) -> dict[str, Any]:
|
||||||
|
meta_path = index_path.parent / f"{index_path.name}.meta.json"
|
||||||
|
if not meta_path.exists():
|
||||||
|
raise FileNotFoundError(f"Metadata file not found: {meta_path}")
|
||||||
|
with open(meta_path, encoding="utf-8") as f:
|
||||||
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
def _autodetect_index_base(index_dir: Path) -> Optional[Path]:
|
||||||
|
"""If exactly one *.leann.meta.json exists, return its base path (without .meta.json)."""
|
||||||
|
candidates = list(index_dir.glob("*.leann.meta.json"))
|
||||||
|
if len(candidates) == 1:
|
||||||
|
meta = candidates[0]
|
||||||
|
base = meta.with_suffix("") # remove .json
|
||||||
|
base = base.with_suffix("") # remove .meta
|
||||||
|
return base
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _load_offset_map(offsets_file: Path) -> dict[str, int]:
|
||||||
|
if not offsets_file.exists():
|
||||||
|
return {}
|
||||||
|
with open(offsets_file, "rb") as f:
|
||||||
|
return pickle.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
def _next_numeric_id(existing_ids: list[str]) -> int:
|
||||||
|
numeric_ids = [int(x) for x in existing_ids if x.isdigit()]
|
||||||
|
if not numeric_ids:
|
||||||
|
return 0
|
||||||
|
return max(numeric_ids) + 1
|
||||||
|
|
||||||
|
|
||||||
|
def build_base_index(
|
||||||
|
base_dir: str,
|
||||||
|
index_dir: str,
|
||||||
|
index_name: str,
|
||||||
|
embedding_model: str,
|
||||||
|
embedding_mode: str,
|
||||||
|
chunk_size: int,
|
||||||
|
chunk_overlap: int,
|
||||||
|
file_types: Optional[list[str]] = None,
|
||||||
|
max_items: int = -1,
|
||||||
|
ef_construction: Optional[int] = None,
|
||||||
|
recompute_build: bool = False,
|
||||||
|
) -> str:
|
||||||
|
print(f"Building base index from: {base_dir}")
|
||||||
|
documents = _load_documents(base_dir, required_exts=file_types)
|
||||||
|
if not documents:
|
||||||
|
raise ValueError(f"No documents found in base_dir: {base_dir}")
|
||||||
|
|
||||||
|
from chunking import create_text_chunks
|
||||||
|
|
||||||
|
texts = create_text_chunks(
|
||||||
|
documents,
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=chunk_overlap,
|
||||||
|
use_ast_chunking=False,
|
||||||
|
)
|
||||||
|
if max_items > 0 and len(texts) > max_items:
|
||||||
|
texts = texts[:max_items]
|
||||||
|
print(f"Limiting to {max_items} chunks")
|
||||||
|
|
||||||
|
index_dir_path = Path(index_dir)
|
||||||
|
_ensure_index_dir(index_dir_path)
|
||||||
|
index_path = index_dir_path / index_name
|
||||||
|
|
||||||
|
print("Creating HNSW index (non-compact)...")
|
||||||
|
from leann.api import LeannBuilder
|
||||||
|
from leann.registry import register_project_directory
|
||||||
|
|
||||||
|
builder = LeannBuilder(
|
||||||
|
backend_name="hnsw",
|
||||||
|
embedding_model=embedding_model,
|
||||||
|
embedding_mode=embedding_mode,
|
||||||
|
is_recompute=recompute_build,
|
||||||
|
is_compact=False,
|
||||||
|
efConstruction=(ef_construction if ef_construction is not None else 200),
|
||||||
|
)
|
||||||
|
for t in texts:
|
||||||
|
builder.add_text(t)
|
||||||
|
builder.build_index(str(index_path))
|
||||||
|
|
||||||
|
# Register for discovery
|
||||||
|
register_project_directory(Path.cwd())
|
||||||
|
|
||||||
|
print(f"Base index built at: {index_path}")
|
||||||
|
return str(index_path)
|
||||||
|
|
||||||
|
|
||||||
|
def add_incremental(
|
||||||
|
add_dir: str,
|
||||||
|
index_dir: str,
|
||||||
|
index_name: Optional[str] = None,
|
||||||
|
embedding_model: Optional[str] = None,
|
||||||
|
embedding_mode: Optional[str] = None,
|
||||||
|
chunk_size: int = 256,
|
||||||
|
chunk_overlap: int = 128,
|
||||||
|
file_types: Optional[list[str]] = None,
|
||||||
|
max_items: int = -1,
|
||||||
|
ef_construction: Optional[int] = None,
|
||||||
|
recompute_add: bool = False,
|
||||||
|
) -> str:
|
||||||
|
print(f"Adding incremental data from: {add_dir}")
|
||||||
|
index_dir_path = Path(index_dir)
|
||||||
|
index_path = index_dir_path / (index_name or "documents.leann")
|
||||||
|
|
||||||
|
# If specified base doesn't exist, try to auto-detect an existing base
|
||||||
|
try:
|
||||||
|
_read_meta(index_path)
|
||||||
|
except FileNotFoundError:
|
||||||
|
auto_base = _autodetect_index_base(index_dir_path)
|
||||||
|
if auto_base is not None:
|
||||||
|
print(f"Auto-detected index base: {auto_base.name}")
|
||||||
|
index_path = auto_base
|
||||||
|
_read_meta(index_path)
|
||||||
|
else:
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"No index metadata found for base '{index_path.name}'. Build base first with --build-base "
|
||||||
|
f"or provide --index-name to match an existing index (e.g., 'test_doc_files.leann')."
|
||||||
|
)
|
||||||
|
|
||||||
|
# Prepare validated context from core (checks backend/no-recompute and resolves embedding defaults)
|
||||||
|
from leann.api import create_incremental_add_context, incremental_add_texts_with_context
|
||||||
|
|
||||||
|
ctx = create_incremental_add_context(
|
||||||
|
str(index_path),
|
||||||
|
embedding_model=embedding_model,
|
||||||
|
embedding_mode=embedding_mode,
|
||||||
|
data_dir=add_dir,
|
||||||
|
required_exts=file_types,
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=chunk_overlap,
|
||||||
|
max_items=max_items,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Use prepared texts from context to perform the add
|
||||||
|
prepared_texts = ctx.prepared_texts or []
|
||||||
|
if not prepared_texts:
|
||||||
|
print("No new chunks to add.")
|
||||||
|
return str(index_path)
|
||||||
|
|
||||||
|
added = incremental_add_texts_with_context(
|
||||||
|
ctx,
|
||||||
|
prepared_texts,
|
||||||
|
ef_construction=ef_construction,
|
||||||
|
recompute=recompute_add,
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f"Incremental add completed. Added {added} chunks. Index: {index_path}")
|
||||||
|
return str(index_path)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description="Dynamic add to LEANN HNSW index without recompute",
|
||||||
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument("--build-base", action="store_true", help="Build base index")
|
||||||
|
parser.add_argument("--add-incremental", action="store_true", help="Add incremental data")
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"--base-dir",
|
||||||
|
type=str,
|
||||||
|
default="/Users/yichuan/Desktop/code/LEANN/leann/data",
|
||||||
|
help="Base data directory",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--add-dir",
|
||||||
|
type=str,
|
||||||
|
default="/Users/yichuan/Desktop/code/LEANN/leann/test_add",
|
||||||
|
help="Incremental data directory",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--index-dir",
|
||||||
|
type=str,
|
||||||
|
default="./test_doc_files",
|
||||||
|
help="Directory containing the index",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--index-name",
|
||||||
|
type=str,
|
||||||
|
default="documents.leann",
|
||||||
|
help=(
|
||||||
|
"Index base file name. If you built via document_rag.py, use 'test_doc_files.leann'. "
|
||||||
|
"Default: documents.leann"
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"--embedding-model",
|
||||||
|
type=str,
|
||||||
|
default="facebook/contriever",
|
||||||
|
help="Embedding model name",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--embedding-mode",
|
||||||
|
type=str,
|
||||||
|
default="sentence-transformers",
|
||||||
|
choices=["sentence-transformers", "openai", "mlx", "ollama"],
|
||||||
|
help="Embedding backend mode",
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument("--chunk-size", type=int, default=256)
|
||||||
|
parser.add_argument("--chunk-overlap", type=int, default=128)
|
||||||
|
parser.add_argument("--file-types", nargs="+", default=None)
|
||||||
|
parser.add_argument("--max-items", type=int, default=-1)
|
||||||
|
parser.add_argument("--ef-construction", type=int, default=32)
|
||||||
|
parser.add_argument(
|
||||||
|
"--recompute-add", action="store_true", help="Enable recompute-mode add (non-compact only)"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--recompute-build",
|
||||||
|
action="store_true",
|
||||||
|
help="Enable recompute-mode base build (non-compact only)",
|
||||||
|
)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
if not args.build_base and not args.add_incremental:
|
||||||
|
print("Nothing to do. Use --build-base and/or --add-incremental.")
|
||||||
|
return
|
||||||
|
|
||||||
|
index_path_str: Optional[str] = None
|
||||||
|
|
||||||
|
if args.build_base:
|
||||||
|
index_path_str = build_base_index(
|
||||||
|
base_dir=args.base_dir,
|
||||||
|
index_dir=args.index_dir,
|
||||||
|
index_name=args.index_name,
|
||||||
|
embedding_model=args.embedding_model,
|
||||||
|
embedding_mode=args.embedding_mode,
|
||||||
|
chunk_size=args.chunk_size,
|
||||||
|
chunk_overlap=args.chunk_overlap,
|
||||||
|
file_types=args.file_types,
|
||||||
|
max_items=args.max_items,
|
||||||
|
ef_construction=args.ef_construction,
|
||||||
|
recompute_build=args.recompute_build,
|
||||||
|
)
|
||||||
|
|
||||||
|
if args.add_incremental:
|
||||||
|
index_path_str = add_incremental(
|
||||||
|
add_dir=args.add_dir,
|
||||||
|
index_dir=args.index_dir,
|
||||||
|
index_name=args.index_name,
|
||||||
|
embedding_model=args.embedding_model,
|
||||||
|
embedding_mode=args.embedding_mode,
|
||||||
|
chunk_size=args.chunk_size,
|
||||||
|
chunk_overlap=args.chunk_overlap,
|
||||||
|
file_types=args.file_types,
|
||||||
|
max_items=args.max_items,
|
||||||
|
ef_construction=args.ef_construction,
|
||||||
|
recompute_add=args.recompute_add,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Optional: quick test query using searcher
|
||||||
|
if index_path_str:
|
||||||
|
try:
|
||||||
|
from leann.api import LeannSearcher
|
||||||
|
|
||||||
|
searcher = LeannSearcher(index_path_str)
|
||||||
|
query = "what is LEANN?"
|
||||||
|
if args.add_incremental:
|
||||||
|
query = "what is the multi vector search and how it works?"
|
||||||
|
results = searcher.search(query, top_k=5)
|
||||||
|
if results:
|
||||||
|
print(f"Sample result: {results[0].text[:80]}...")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -1,5 +1,5 @@
|
|||||||
[build-system]
|
[build-system]
|
||||||
requires = ["scikit-build-core>=0.10", "pybind11>=2.12.0", "numpy", "cmake>=3.30"]
|
requires = ["scikit-build-core>=0.10", "pybind11>=2.12.0", "numpy"]
|
||||||
build-backend = "scikit_build_core.build"
|
build-backend = "scikit_build_core.build"
|
||||||
|
|
||||||
[project]
|
[project]
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ from leann.registry import register_backend
|
|||||||
from leann.searcher_base import BaseSearcher
|
from leann.searcher_base import BaseSearcher
|
||||||
|
|
||||||
from .convert_to_csr import convert_hnsw_graph_to_csr
|
from .convert_to_csr import convert_hnsw_graph_to_csr
|
||||||
|
from .prune_index import prune_embeddings_preserve_graph_inplace
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -90,8 +91,16 @@ class HNSWBuilder(LeannBackendBuilderInterface):
|
|||||||
index_file = index_dir / f"{index_prefix}.index"
|
index_file = index_dir / f"{index_prefix}.index"
|
||||||
faiss.write_index(index, str(index_file))
|
faiss.write_index(index, str(index_file))
|
||||||
|
|
||||||
if self.is_compact:
|
if self.is_recompute:
|
||||||
self._convert_to_csr(index_file)
|
if self.is_compact:
|
||||||
|
self._convert_to_csr(index_file)
|
||||||
|
else:
|
||||||
|
# Non-compact format: prune only embeddings, keep original graph
|
||||||
|
ok = prune_embeddings_preserve_graph_inplace(str(index_file))
|
||||||
|
if not ok:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Pruning embeddings while preserving graph failed for non-compact index"
|
||||||
|
)
|
||||||
|
|
||||||
def _convert_to_csr(self, index_file: Path):
|
def _convert_to_csr(self, index_file: Path):
|
||||||
"""Convert built index to CSR format"""
|
"""Convert built index to CSR format"""
|
||||||
@@ -148,7 +157,13 @@ class HNSWSearcher(BaseSearcher):
|
|||||||
self.is_pruned
|
self.is_pruned
|
||||||
) # In C++ code, it's called is_recompute, but it's only for loading IIUC.
|
) # In C++ code, it's called is_recompute, but it's only for loading IIUC.
|
||||||
|
|
||||||
self._index = faiss.read_index(str(index_file), faiss.IO_FLAG_MMAP, hnsw_config)
|
# If pruned (recompute mode), explicitly skip storage to avoid reading
|
||||||
|
# the pruned section. Still allow MMAP for graph.
|
||||||
|
io_flags = faiss.IO_FLAG_MMAP
|
||||||
|
if self.is_pruned:
|
||||||
|
io_flags |= faiss.IO_FLAG_SKIP_STORAGE
|
||||||
|
|
||||||
|
self._index = faiss.read_index(str(index_file), io_flags, hnsw_config)
|
||||||
|
|
||||||
def search(
|
def search(
|
||||||
self,
|
self,
|
||||||
@@ -251,3 +266,55 @@ class HNSWSearcher(BaseSearcher):
|
|||||||
string_labels = [[str(int_label) for int_label in batch_labels] for batch_labels in labels]
|
string_labels = [[str(int_label) for int_label in batch_labels] for batch_labels in labels]
|
||||||
|
|
||||||
return {"labels": string_labels, "distances": distances}
|
return {"labels": string_labels, "distances": distances}
|
||||||
|
|
||||||
|
|
||||||
|
# ---------- Helper API for incremental add (Python-level) ----------
|
||||||
|
def add_vectors(
|
||||||
|
index_file_path: str,
|
||||||
|
embeddings: np.ndarray,
|
||||||
|
*,
|
||||||
|
ef_construction: Optional[int] = None,
|
||||||
|
recompute: bool = False,
|
||||||
|
) -> None:
|
||||||
|
"""Append vectors to an existing non-compact HNSW index.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
index_file_path: Path to the HNSW .index file
|
||||||
|
embeddings: float32 numpy array (N, D)
|
||||||
|
ef_construction: Optional override for efConstruction during insertion
|
||||||
|
recompute: Reserved for future use to control insertion-time recompute behaviors
|
||||||
|
"""
|
||||||
|
from . import faiss # type: ignore
|
||||||
|
|
||||||
|
if embeddings.dtype != np.float32:
|
||||||
|
embeddings = embeddings.astype(np.float32)
|
||||||
|
if not embeddings.flags.c_contiguous:
|
||||||
|
embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
|
||||||
|
|
||||||
|
# Load index normally to ensure storage is present; toggle is_recompute on the object
|
||||||
|
index = faiss.read_index(str(index_file_path), faiss.IO_FLAG_MMAP)
|
||||||
|
|
||||||
|
# Best-effort: explicitly set flag on the object if the binding exposes it
|
||||||
|
try:
|
||||||
|
index.is_recompute = bool(recompute)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
if ef_construction is not None:
|
||||||
|
index.hnsw.efConstruction = int(ef_construction)
|
||||||
|
except Exception:
|
||||||
|
# Best-effort; ignore if backend doesn't expose setter
|
||||||
|
pass
|
||||||
|
|
||||||
|
# For non-compact HNSW, calling add directly is sufficient. When is_recompute is set
|
||||||
|
# (via config or attribute), FAISS will run the insertion/search path accordingly.
|
||||||
|
# To strictly follow per-point insert semantics in recompute mode, add one-by-one.
|
||||||
|
if recompute:
|
||||||
|
# Insert row by row
|
||||||
|
n = embeddings.shape[0]
|
||||||
|
for i in range(n):
|
||||||
|
row = embeddings[i : i + 1]
|
||||||
|
index.add(1, faiss.swig_ptr(row))
|
||||||
|
else:
|
||||||
|
index.add(embeddings.shape[0], faiss.swig_ptr(embeddings))
|
||||||
|
faiss.write_index(index, str(index_file_path))
|
||||||
|
|||||||
149
packages/leann-backend-hnsw/leann_backend_hnsw/prune_index.py
Normal file
149
packages/leann-backend-hnsw/leann_backend_hnsw/prune_index.py
Normal file
@@ -0,0 +1,149 @@
|
|||||||
|
import os
|
||||||
|
import struct
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from .convert_to_csr import (
|
||||||
|
EXPECTED_HNSW_FOURCCS,
|
||||||
|
NULL_INDEX_FOURCC,
|
||||||
|
read_struct,
|
||||||
|
read_vector_raw,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _write_vector_raw(f_out, count: int, data_bytes: bytes) -> None:
|
||||||
|
"""Write a vector in the same binary layout as read_vector_raw reads: <Q count> + raw bytes."""
|
||||||
|
f_out.write(struct.pack("<Q", count))
|
||||||
|
if count > 0 and data_bytes:
|
||||||
|
f_out.write(data_bytes)
|
||||||
|
|
||||||
|
|
||||||
|
def prune_embeddings_preserve_graph(input_filename: str, output_filename: str) -> bool:
|
||||||
|
"""
|
||||||
|
Copy an original (non-compact) HNSW index file while pruning the trailing embedding storage.
|
||||||
|
Preserves the graph structure and metadata exactly; only writes a NULL storage marker instead of
|
||||||
|
the original storage fourcc and payload.
|
||||||
|
|
||||||
|
Returns True on success.
|
||||||
|
"""
|
||||||
|
print(f"Pruning embeddings from {input_filename} to {output_filename}")
|
||||||
|
print("--------------------------------")
|
||||||
|
# running in mode is-recompute=True and is-compact=False
|
||||||
|
in_path = Path(input_filename)
|
||||||
|
out_path = Path(output_filename)
|
||||||
|
|
||||||
|
try:
|
||||||
|
with open(in_path, "rb") as f_in, open(out_path, "wb") as f_out:
|
||||||
|
# Header
|
||||||
|
index_fourcc = read_struct(f_in, "<I")
|
||||||
|
if index_fourcc not in EXPECTED_HNSW_FOURCCS:
|
||||||
|
# Still proceed, but this is unexpected
|
||||||
|
pass
|
||||||
|
f_out.write(struct.pack("<I", index_fourcc))
|
||||||
|
|
||||||
|
d = read_struct(f_in, "<i")
|
||||||
|
ntotal_hdr = read_struct(f_in, "<q")
|
||||||
|
dummy1 = read_struct(f_in, "<q")
|
||||||
|
dummy2 = read_struct(f_in, "<q")
|
||||||
|
is_trained = read_struct(f_in, "?")
|
||||||
|
metric_type = read_struct(f_in, "<i")
|
||||||
|
f_out.write(struct.pack("<i", d))
|
||||||
|
f_out.write(struct.pack("<q", ntotal_hdr))
|
||||||
|
f_out.write(struct.pack("<q", dummy1))
|
||||||
|
f_out.write(struct.pack("<q", dummy2))
|
||||||
|
f_out.write(struct.pack("<?", is_trained))
|
||||||
|
f_out.write(struct.pack("<i", metric_type))
|
||||||
|
|
||||||
|
if metric_type > 1:
|
||||||
|
metric_arg = read_struct(f_in, "<f")
|
||||||
|
f_out.write(struct.pack("<f", metric_arg))
|
||||||
|
|
||||||
|
# Vectors: assign_probas (double), cum_nneighbor_per_level (int32), levels (int32)
|
||||||
|
cnt, data = read_vector_raw(f_in, "d")
|
||||||
|
_write_vector_raw(f_out, cnt, data)
|
||||||
|
|
||||||
|
cnt, data = read_vector_raw(f_in, "i")
|
||||||
|
_write_vector_raw(f_out, cnt, data)
|
||||||
|
|
||||||
|
cnt, data = read_vector_raw(f_in, "i")
|
||||||
|
_write_vector_raw(f_out, cnt, data)
|
||||||
|
|
||||||
|
# Probe potential extra alignment/flag byte present in some original formats
|
||||||
|
probe = f_in.read(1)
|
||||||
|
if probe:
|
||||||
|
if probe == b"\x00":
|
||||||
|
# Preserve this unexpected 0x00 byte
|
||||||
|
f_out.write(probe)
|
||||||
|
else:
|
||||||
|
# Likely part of the next vector; rewind
|
||||||
|
f_in.seek(-1, os.SEEK_CUR)
|
||||||
|
|
||||||
|
# Offsets (uint64) and neighbors (int32)
|
||||||
|
cnt, data = read_vector_raw(f_in, "Q")
|
||||||
|
_write_vector_raw(f_out, cnt, data)
|
||||||
|
|
||||||
|
cnt, data = read_vector_raw(f_in, "i")
|
||||||
|
_write_vector_raw(f_out, cnt, data)
|
||||||
|
|
||||||
|
# Scalar params
|
||||||
|
entry_point = read_struct(f_in, "<i")
|
||||||
|
max_level = read_struct(f_in, "<i")
|
||||||
|
ef_construction = read_struct(f_in, "<i")
|
||||||
|
ef_search = read_struct(f_in, "<i")
|
||||||
|
dummy_upper_beam = read_struct(f_in, "<i")
|
||||||
|
f_out.write(struct.pack("<i", entry_point))
|
||||||
|
f_out.write(struct.pack("<i", max_level))
|
||||||
|
f_out.write(struct.pack("<i", ef_construction))
|
||||||
|
f_out.write(struct.pack("<i", ef_search))
|
||||||
|
f_out.write(struct.pack("<i", dummy_upper_beam))
|
||||||
|
|
||||||
|
# Storage fourcc (if present) — write NULL marker and drop any remaining data
|
||||||
|
try:
|
||||||
|
read_struct(f_in, "<I")
|
||||||
|
# Regardless of original, write NULL
|
||||||
|
f_out.write(struct.pack("<I", NULL_INDEX_FOURCC))
|
||||||
|
# Discard the rest of the file (embedding payload)
|
||||||
|
# (Do not copy anything else)
|
||||||
|
except EOFError:
|
||||||
|
# No storage section; nothing else to write
|
||||||
|
pass
|
||||||
|
|
||||||
|
return True
|
||||||
|
except Exception:
|
||||||
|
# Best-effort cleanup
|
||||||
|
try:
|
||||||
|
if out_path.exists():
|
||||||
|
out_path.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def prune_embeddings_preserve_graph_inplace(index_file_path: str) -> bool:
|
||||||
|
"""
|
||||||
|
Convenience wrapper: write pruned file to a temporary path next to the
|
||||||
|
original, then atomically replace on success.
|
||||||
|
"""
|
||||||
|
print(f"Pruning embeddings from {index_file_path} to {index_file_path}")
|
||||||
|
print("--------------------------------")
|
||||||
|
# running in mode is-recompute=True and is-compact=False
|
||||||
|
src = Path(index_file_path)
|
||||||
|
tmp = src.with_suffix(".pruned.tmp")
|
||||||
|
ok = prune_embeddings_preserve_graph(str(src), str(tmp))
|
||||||
|
if not ok:
|
||||||
|
if tmp.exists():
|
||||||
|
try:
|
||||||
|
tmp.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
return False
|
||||||
|
try:
|
||||||
|
os.replace(str(tmp), str(src))
|
||||||
|
except Exception:
|
||||||
|
# Rollback on failure
|
||||||
|
try:
|
||||||
|
if tmp.exists():
|
||||||
|
tmp.unlink()
|
||||||
|
except OSError:
|
||||||
|
pass
|
||||||
|
return False
|
||||||
|
return True
|
||||||
Submodule packages/leann-backend-hnsw/third_party/faiss updated: ed96ff7dba...ea86d06ceb
@@ -5,6 +5,7 @@ with the correct, original embedding logic from the user's reference code.
|
|||||||
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
import pickle
|
import pickle
|
||||||
import re
|
import re
|
||||||
import subprocess
|
import subprocess
|
||||||
@@ -19,6 +20,7 @@ import numpy as np
|
|||||||
from leann.interface import LeannBackendSearcherInterface
|
from leann.interface import LeannBackendSearcherInterface
|
||||||
|
|
||||||
from .chat import get_llm
|
from .chat import get_llm
|
||||||
|
from .embedding_server_manager import EmbeddingServerManager
|
||||||
from .interface import LeannBackendFactoryInterface
|
from .interface import LeannBackendFactoryInterface
|
||||||
from .metadata_filter import MetadataFilterEngine
|
from .metadata_filter import MetadataFilterEngine
|
||||||
from .registry import BACKEND_REGISTRY
|
from .registry import BACKEND_REGISTRY
|
||||||
@@ -118,6 +120,20 @@ class SearchResult:
|
|||||||
metadata: dict[str, Any] = field(default_factory=dict)
|
metadata: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class IncrementalAddContext:
|
||||||
|
"""Prepared context for safe incremental add operations on an index."""
|
||||||
|
|
||||||
|
index_path: str
|
||||||
|
passages_file: Path
|
||||||
|
offsets_file: Path
|
||||||
|
vector_index_file: Path
|
||||||
|
embedding_model: str
|
||||||
|
embedding_mode: str
|
||||||
|
distance_metric: str
|
||||||
|
prepared_texts: Optional[list[str]] = None
|
||||||
|
|
||||||
|
|
||||||
class PassageManager:
|
class PassageManager:
|
||||||
def __init__(
|
def __init__(
|
||||||
self, passage_sources: list[dict[str, Any]], metadata_file_path: Optional[str] = None
|
self, passage_sources: list[dict[str, Any]], metadata_file_path: Optional[str] = None
|
||||||
@@ -476,9 +492,7 @@ class LeannBuilder:
|
|||||||
is_compact = self.backend_kwargs.get("is_compact", True)
|
is_compact = self.backend_kwargs.get("is_compact", True)
|
||||||
is_recompute = self.backend_kwargs.get("is_recompute", True)
|
is_recompute = self.backend_kwargs.get("is_recompute", True)
|
||||||
meta_data["is_compact"] = is_compact
|
meta_data["is_compact"] = is_compact
|
||||||
meta_data["is_pruned"] = (
|
meta_data["is_pruned"] = is_recompute # Pruned only if compact and recompute
|
||||||
is_compact and is_recompute
|
|
||||||
) # Pruned only if compact and recompute
|
|
||||||
with open(leann_meta_path, "w", encoding="utf-8") as f:
|
with open(leann_meta_path, "w", encoding="utf-8") as f:
|
||||||
json.dump(meta_data, f, indent=2)
|
json.dump(meta_data, f, indent=2)
|
||||||
|
|
||||||
@@ -1018,8 +1032,405 @@ class LeannChat:
|
|||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
def __del__(self):
|
|
||||||
|
# ------------------------------
|
||||||
|
# Incremental Add Utilities (HNSW no-recompute only)
|
||||||
|
# ------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_index_paths(index_path: str) -> tuple[Path, Path, Path]:
|
||||||
|
"""Given base index path (without extension), return (passages.jsonl, passages.idx, vector.index).
|
||||||
|
|
||||||
|
For HNSW, vector index file is typically <stem>.index (e.g., documents.index) even when base is
|
||||||
|
'documents.leann'. We prefer an existing <stem>.index, otherwise fall back to <name>.index.
|
||||||
|
"""
|
||||||
|
base = Path(index_path)
|
||||||
|
passages_file = base.parent / f"{base.name}.passages.jsonl"
|
||||||
|
offsets_file = base.parent / f"{base.name}.passages.idx"
|
||||||
|
candidate_name_index = base.parent / f"{base.name}.index"
|
||||||
|
candidate_stem_index = base.parent / f"{base.stem}.index"
|
||||||
|
vector_index_file = (
|
||||||
|
candidate_stem_index if candidate_stem_index.exists() else candidate_name_index
|
||||||
|
)
|
||||||
|
return passages_file, offsets_file, vector_index_file
|
||||||
|
|
||||||
|
|
||||||
|
def _read_meta_file(index_path: str) -> dict[str, Any]:
|
||||||
|
meta_path = Path(f"{index_path}.meta.json")
|
||||||
|
if not meta_path.exists():
|
||||||
|
raise FileNotFoundError(f"Leann metadata file not found: {meta_path}")
|
||||||
|
with open(meta_path, encoding="utf-8") as f:
|
||||||
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
def _load_offset_map_pickle(offsets_file: Path) -> dict[str, int]:
|
||||||
|
if not offsets_file.exists():
|
||||||
|
return {}
|
||||||
|
with open(offsets_file, "rb") as f:
|
||||||
|
return pickle.load(f)
|
||||||
|
|
||||||
|
|
||||||
|
def _append_passages_and_update_offsets(
|
||||||
|
passages_file: Path, offsets_file: Path, new_texts: list[str]
|
||||||
|
) -> list[str]:
|
||||||
|
"""Append new texts to passages file, update offset map, and return assigned string IDs.
|
||||||
|
|
||||||
|
IDs are assigned as incrementing integers based on existing keys in the offset map.
|
||||||
|
"""
|
||||||
|
offset_map = _load_offset_map_pickle(offsets_file)
|
||||||
|
# Compute next numeric id
|
||||||
|
numeric_ids = [int(x) for x in offset_map.keys() if str(x).isdigit()]
|
||||||
|
next_id_num = (max(numeric_ids) + 1) if numeric_ids else 0
|
||||||
|
assigned_ids: list[str] = []
|
||||||
|
|
||||||
|
with open(passages_file, "a", encoding="utf-8") as f:
|
||||||
|
for text in new_texts:
|
||||||
|
offset = f.tell()
|
||||||
|
str_id = str(next_id_num)
|
||||||
|
json.dump({"id": str_id, "text": text, "metadata": {}}, f, ensure_ascii=False)
|
||||||
|
f.write("\n")
|
||||||
|
offset_map[str_id] = offset
|
||||||
|
assigned_ids.append(str_id)
|
||||||
|
next_id_num += 1
|
||||||
|
|
||||||
|
with open(offsets_file, "wb") as f:
|
||||||
|
pickle.dump(offset_map, f)
|
||||||
|
|
||||||
|
return assigned_ids
|
||||||
|
|
||||||
|
|
||||||
|
def incremental_add_texts(
|
||||||
|
index_path: str,
|
||||||
|
texts: list[str],
|
||||||
|
*,
|
||||||
|
embedding_model: Optional[str] = None,
|
||||||
|
embedding_mode: Optional[str] = None,
|
||||||
|
ef_construction: Optional[int] = None,
|
||||||
|
recompute: bool = False,
|
||||||
|
) -> int:
|
||||||
|
"""Incrementally add text chunks to an existing HNSW index built with no-recompute.
|
||||||
|
|
||||||
|
- Validates backend is HNSW and index is non-compact (no-recompute path)
|
||||||
|
- Appends passages and offsets
|
||||||
|
- Computes embeddings and appends to the HNSW vector index
|
||||||
|
|
||||||
|
Returns number of added chunks.
|
||||||
|
"""
|
||||||
|
if not texts:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
meta = _read_meta_file(index_path)
|
||||||
|
if meta.get("backend_name") != "hnsw":
|
||||||
|
raise RuntimeError("Incremental add is currently supported only for HNSW backend")
|
||||||
|
if meta.get("is_compact", True):
|
||||||
|
raise RuntimeError(
|
||||||
|
"Index is compact/pruned. Rebuild base with is_recompute=False and is_compact=False for incremental add."
|
||||||
|
)
|
||||||
|
|
||||||
|
passages_file, offsets_file, vector_index_file = _resolve_index_paths(index_path)
|
||||||
|
if not vector_index_file.exists():
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"Vector index file missing: {vector_index_file}. Build base first with LeannBuilder."
|
||||||
|
)
|
||||||
|
|
||||||
|
# Resolve embedding config from meta if not provided
|
||||||
|
model_name = embedding_model or meta.get("embedding_model", "facebook/contriever")
|
||||||
|
mode_name = embedding_mode or meta.get("embedding_mode", "sentence-transformers")
|
||||||
|
|
||||||
|
# Append passages and update offsets
|
||||||
|
assigned_ids = _append_passages_and_update_offsets(passages_file, offsets_file, texts)
|
||||||
|
|
||||||
|
# Compute embeddings
|
||||||
|
# Embedding computation path
|
||||||
|
esm = None
|
||||||
|
port = None
|
||||||
|
if recompute:
|
||||||
|
# Determine distance metric early for server config
|
||||||
|
distance_metric = meta.get("backend_kwargs", {}).get("distance_metric", "mips").lower()
|
||||||
|
# Start embedding server and compute via ZMQ for consistency with recompute semantics
|
||||||
|
passages_source_file = f"{index_path}.meta.json"
|
||||||
|
esm = EmbeddingServerManager(
|
||||||
|
backend_module_name="leann_backend_hnsw.hnsw_embedding_server",
|
||||||
|
)
|
||||||
|
started, port = esm.start_server(
|
||||||
|
port=5557,
|
||||||
|
model_name=model_name,
|
||||||
|
embedding_mode=mode_name,
|
||||||
|
passages_file=passages_source_file,
|
||||||
|
distance_metric=distance_metric,
|
||||||
|
enable_warmup=False,
|
||||||
|
)
|
||||||
|
if not started:
|
||||||
|
raise RuntimeError("Failed to start embedding server for recompute add")
|
||||||
|
embeddings = compute_embeddings_via_server(texts, model_name, port)
|
||||||
|
else:
|
||||||
|
embeddings = compute_embeddings(
|
||||||
|
texts,
|
||||||
|
model_name=model_name,
|
||||||
|
mode=mode_name,
|
||||||
|
use_server=False,
|
||||||
|
is_build=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Normalize for cosine if needed
|
||||||
|
if "distance_metric" not in locals():
|
||||||
|
distance_metric = meta.get("backend_kwargs", {}).get("distance_metric", "mips").lower()
|
||||||
|
if distance_metric == "cosine":
|
||||||
|
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||||
|
norms[norms == 0] = 1
|
||||||
|
embeddings = embeddings / norms
|
||||||
|
|
||||||
|
# Append via backend helper (supports ef_construction/recompute plumbing)
|
||||||
|
try:
|
||||||
|
from leann_backend_hnsw.hnsw_backend import add_vectors as hnsw_add_vectors # type: ignore
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Failed to import HNSW backend add helper. Ensure HNSW backend is installed."
|
||||||
|
) from e
|
||||||
|
|
||||||
|
# Propagate ZMQ port to FAISS add path when recompute is True
|
||||||
|
if recompute and port is not None:
|
||||||
|
os.environ["LEANN_ZMQ_PORT"] = str(port)
|
||||||
|
|
||||||
|
hnsw_add_vectors(
|
||||||
|
str(vector_index_file),
|
||||||
|
embeddings,
|
||||||
|
ef_construction=ef_construction,
|
||||||
|
recompute=recompute,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Stop server after add when recompute path used
|
||||||
|
if esm is not None:
|
||||||
try:
|
try:
|
||||||
self.cleanup()
|
esm.stop_server()
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
# Sanity: ids length should match embeddings rows
|
||||||
|
if len(assigned_ids) != embeddings.shape[0]:
|
||||||
|
warnings.warn(
|
||||||
|
f"Assigned {len(assigned_ids)} IDs but computed {embeddings.shape[0]} embeddings.",
|
||||||
|
UserWarning,
|
||||||
|
stacklevel=2,
|
||||||
|
)
|
||||||
|
|
||||||
|
return len(assigned_ids)
|
||||||
|
|
||||||
|
|
||||||
|
def create_incremental_add_context(
|
||||||
|
index_path: str,
|
||||||
|
*,
|
||||||
|
# Optional embedding choices; if None will use meta
|
||||||
|
embedding_model: Optional[str] = None,
|
||||||
|
embedding_mode: Optional[str] = None,
|
||||||
|
# Optional data-to-text preparation in context
|
||||||
|
data_dir: Optional[str] = None,
|
||||||
|
required_exts: Optional[list[str]] = None,
|
||||||
|
chunk_size: int = 256,
|
||||||
|
chunk_overlap: int = 128,
|
||||||
|
max_items: int = -1,
|
||||||
|
) -> IncrementalAddContext:
|
||||||
|
"""Validate index and prepare context for repeated incremental adds.
|
||||||
|
|
||||||
|
Additionally, if data_dir is provided, this function will load documents,
|
||||||
|
chunk them to texts with the specified parameters, and store them in ctx.prepared_texts.
|
||||||
|
"""
|
||||||
|
meta = _read_meta_file(index_path)
|
||||||
|
if meta.get("backend_name") != "hnsw":
|
||||||
|
raise RuntimeError("Incremental add is currently supported only for HNSW backend")
|
||||||
|
if meta.get("is_compact", True):
|
||||||
|
raise RuntimeError(
|
||||||
|
"Index is compact/pruned. Rebuild base with is_recompute=False and is_compact=False for incremental add."
|
||||||
|
)
|
||||||
|
|
||||||
|
passages_file, offsets_file, vector_index_file = _resolve_index_paths(index_path)
|
||||||
|
if not vector_index_file.exists():
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"Vector index file missing: {vector_index_file}. Build base first with LeannBuilder."
|
||||||
|
)
|
||||||
|
|
||||||
|
model_name = embedding_model or meta.get("embedding_model", "facebook/contriever")
|
||||||
|
mode_name = embedding_mode or meta.get("embedding_mode", "sentence-transformers")
|
||||||
|
distance_metric = meta.get("backend_kwargs", {}).get("distance_metric", "mips").lower()
|
||||||
|
|
||||||
|
prepared_texts: Optional[list[str]] = None
|
||||||
|
if data_dir is not None:
|
||||||
|
try:
|
||||||
|
from llama_index.core import SimpleDirectoryReader # type: ignore
|
||||||
|
from llama_index.core.node_parser import SentenceSplitter # type: ignore
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
"llama-index-core is required when using data_dir in create_incremental_add_context"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
reader_kwargs: dict[str, Any] = {"recursive": True, "encoding": "utf-8"}
|
||||||
|
if required_exts:
|
||||||
|
reader_kwargs["required_exts"] = required_exts
|
||||||
|
documents = SimpleDirectoryReader(data_dir, **reader_kwargs).load_data(show_progress=True)
|
||||||
|
if documents:
|
||||||
|
splitter = SentenceSplitter(
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=chunk_overlap,
|
||||||
|
separator=" ",
|
||||||
|
paragraph_separator="\n\n",
|
||||||
|
)
|
||||||
|
prepared_texts = []
|
||||||
|
for doc in documents:
|
||||||
|
try:
|
||||||
|
nodes = splitter.get_nodes_from_documents([doc])
|
||||||
|
if nodes:
|
||||||
|
prepared_texts.extend([node.get_content() for node in nodes])
|
||||||
|
except Exception:
|
||||||
|
content = doc.get_content()
|
||||||
|
if content and content.strip():
|
||||||
|
prepared_texts.append(content.strip())
|
||||||
|
if max_items > 0 and len(prepared_texts) > max_items:
|
||||||
|
prepared_texts = prepared_texts[:max_items]
|
||||||
|
|
||||||
|
return IncrementalAddContext(
|
||||||
|
index_path=index_path,
|
||||||
|
passages_file=passages_file,
|
||||||
|
offsets_file=offsets_file,
|
||||||
|
vector_index_file=vector_index_file,
|
||||||
|
embedding_model=model_name,
|
||||||
|
embedding_mode=mode_name,
|
||||||
|
distance_metric=distance_metric,
|
||||||
|
prepared_texts=prepared_texts,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def incremental_add_texts_with_context(
|
||||||
|
ctx: IncrementalAddContext,
|
||||||
|
texts: list[str],
|
||||||
|
*,
|
||||||
|
ef_construction: Optional[int] = None,
|
||||||
|
recompute: bool = False,
|
||||||
|
) -> int:
|
||||||
|
"""Incrementally add texts using a prepared context (no repeated validation).
|
||||||
|
|
||||||
|
For non-compact HNSW, ef_construction (efConstruction) can be overridden during insertion.
|
||||||
|
"""
|
||||||
|
if not texts:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
# Append passages & offsets
|
||||||
|
_append_passages_and_update_offsets(ctx.passages_file, ctx.offsets_file, texts)
|
||||||
|
|
||||||
|
# Compute embeddings
|
||||||
|
# Embedding computation path
|
||||||
|
esm = None
|
||||||
|
port = None
|
||||||
|
if recompute:
|
||||||
|
passages_source_file = f"{ctx.index_path}.meta.json"
|
||||||
|
esm = EmbeddingServerManager(
|
||||||
|
backend_module_name="leann_backend_hnsw.hnsw_embedding_server",
|
||||||
|
)
|
||||||
|
started, port = esm.start_server(
|
||||||
|
port=5557,
|
||||||
|
model_name=ctx.embedding_model,
|
||||||
|
embedding_mode=ctx.embedding_mode,
|
||||||
|
passages_file=passages_source_file,
|
||||||
|
distance_metric=ctx.distance_metric,
|
||||||
|
enable_warmup=False,
|
||||||
|
)
|
||||||
|
if not started:
|
||||||
|
raise RuntimeError("Failed to start embedding server for recompute add")
|
||||||
|
embeddings = compute_embeddings_via_server(texts, ctx.embedding_model, port)
|
||||||
|
else:
|
||||||
|
embeddings = compute_embeddings(
|
||||||
|
texts,
|
||||||
|
model_name=ctx.embedding_model,
|
||||||
|
mode=ctx.embedding_mode,
|
||||||
|
use_server=False,
|
||||||
|
is_build=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Normalize for cosine if needed
|
||||||
|
if ctx.distance_metric == "cosine":
|
||||||
|
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||||
|
norms[norms == 0] = 1
|
||||||
|
embeddings = embeddings / norms
|
||||||
|
|
||||||
|
# Append via backend helper (supports ef_construction/recompute plumbing)
|
||||||
|
try:
|
||||||
|
from leann_backend_hnsw.hnsw_backend import add_vectors as hnsw_add_vectors # type: ignore
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError(
|
||||||
|
"Failed to import HNSW backend add helper. Ensure HNSW backend is installed."
|
||||||
|
) from e
|
||||||
|
|
||||||
|
if recompute and port is not None:
|
||||||
|
os.environ["LEANN_ZMQ_PORT"] = str(port)
|
||||||
|
|
||||||
|
hnsw_add_vectors(
|
||||||
|
str(ctx.vector_index_file),
|
||||||
|
embeddings,
|
||||||
|
ef_construction=ef_construction,
|
||||||
|
recompute=recompute,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Stop server after add when recompute path used
|
||||||
|
if esm is not None:
|
||||||
|
try:
|
||||||
|
esm.stop_server()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return embeddings.shape[0]
|
||||||
|
|
||||||
|
|
||||||
|
def incremental_add_directory(
|
||||||
|
index_path: str,
|
||||||
|
data_dir: str,
|
||||||
|
*,
|
||||||
|
chunk_size: int = 256,
|
||||||
|
chunk_overlap: int = 128,
|
||||||
|
required_exts: Optional[list[str]] = None,
|
||||||
|
max_items: int = -1,
|
||||||
|
embedding_model: Optional[str] = None,
|
||||||
|
embedding_mode: Optional[str] = None,
|
||||||
|
) -> int:
|
||||||
|
"""Load documents from a directory, chunk them, and incrementally add to an index.
|
||||||
|
|
||||||
|
Chunking uses LlamaIndex SentenceSplitter for simplicity and avoids external app dependencies.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from llama_index.core import SimpleDirectoryReader # type: ignore
|
||||||
|
from llama_index.core.node_parser import SentenceSplitter # type: ignore
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError("llama-index-core is required for incremental_add_directory") from e
|
||||||
|
|
||||||
|
reader_kwargs: dict[str, Any] = {"recursive": True, "encoding": "utf-8"}
|
||||||
|
if required_exts:
|
||||||
|
reader_kwargs["required_exts"] = required_exts
|
||||||
|
documents = SimpleDirectoryReader(data_dir, **reader_kwargs).load_data(show_progress=True)
|
||||||
|
if not documents:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
# Traditional text chunking
|
||||||
|
splitter = SentenceSplitter(
|
||||||
|
chunk_size=chunk_size,
|
||||||
|
chunk_overlap=chunk_overlap,
|
||||||
|
separator=" ",
|
||||||
|
paragraph_separator="\n\n",
|
||||||
|
)
|
||||||
|
all_texts: list[str] = []
|
||||||
|
for doc in documents:
|
||||||
|
try:
|
||||||
|
nodes = splitter.get_nodes_from_documents([doc])
|
||||||
|
if nodes:
|
||||||
|
all_texts.extend([node.get_content() for node in nodes])
|
||||||
|
except Exception:
|
||||||
|
content = doc.get_content()
|
||||||
|
if content and content.strip():
|
||||||
|
all_texts.append(content.strip())
|
||||||
|
|
||||||
|
if max_items > 0 and len(all_texts) > max_items:
|
||||||
|
all_texts = all_texts[:max_items]
|
||||||
|
|
||||||
|
return incremental_add_texts(
|
||||||
|
index_path,
|
||||||
|
all_texts,
|
||||||
|
embedding_model=embedding_model,
|
||||||
|
embedding_mode=embedding_mode,
|
||||||
|
)
|
||||||
|
|||||||
@@ -104,11 +104,7 @@ astchunk = { path = "packages/astchunk-leann", editable = true }
|
|||||||
[tool.ruff]
|
[tool.ruff]
|
||||||
target-version = "py39"
|
target-version = "py39"
|
||||||
line-length = 100
|
line-length = 100
|
||||||
extend-exclude = [
|
extend-exclude = ["third_party"]
|
||||||
"third_party",
|
|
||||||
"apps/multimodal/vision-based-pdf-multi-vector/multi-vector-leann.py",
|
|
||||||
"apps/multimodal/vision-based-pdf-multi-vector/multi-vector-leann-similarity-map.py"
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
[tool.ruff.lint]
|
[tool.ruff.lint]
|
||||||
|
|||||||
12
uv.lock
generated
12
uv.lock
generated
@@ -1,5 +1,5 @@
|
|||||||
version = 1
|
version = 1
|
||||||
revision = 3
|
revision = 2
|
||||||
requires-python = ">=3.9"
|
requires-python = ">=3.9"
|
||||||
resolution-markers = [
|
resolution-markers = [
|
||||||
"python_full_version >= '3.12'",
|
"python_full_version >= '3.12'",
|
||||||
@@ -2138,7 +2138,7 @@ wheels = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "leann-backend-diskann"
|
name = "leann-backend-diskann"
|
||||||
version = "0.3.4"
|
version = "0.3.3"
|
||||||
source = { editable = "packages/leann-backend-diskann" }
|
source = { editable = "packages/leann-backend-diskann" }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "leann-core" },
|
{ name = "leann-core" },
|
||||||
@@ -2150,14 +2150,14 @@ dependencies = [
|
|||||||
|
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
{ name = "leann-core", specifier = "==0.3.4" },
|
{ name = "leann-core", specifier = "==0.3.3" },
|
||||||
{ name = "numpy" },
|
{ name = "numpy" },
|
||||||
{ name = "protobuf", specifier = ">=3.19.0" },
|
{ name = "protobuf", specifier = ">=3.19.0" },
|
||||||
]
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "leann-backend-hnsw"
|
name = "leann-backend-hnsw"
|
||||||
version = "0.3.4"
|
version = "0.3.3"
|
||||||
source = { editable = "packages/leann-backend-hnsw" }
|
source = { editable = "packages/leann-backend-hnsw" }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "leann-core" },
|
{ name = "leann-core" },
|
||||||
@@ -2170,7 +2170,7 @@ dependencies = [
|
|||||||
|
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
{ name = "leann-core", specifier = "==0.3.4" },
|
{ name = "leann-core", specifier = "==0.3.3" },
|
||||||
{ name = "msgpack", specifier = ">=1.0.0" },
|
{ name = "msgpack", specifier = ">=1.0.0" },
|
||||||
{ name = "numpy" },
|
{ name = "numpy" },
|
||||||
{ name = "pyzmq", specifier = ">=23.0.0" },
|
{ name = "pyzmq", specifier = ">=23.0.0" },
|
||||||
@@ -2178,7 +2178,7 @@ requires-dist = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "leann-core"
|
name = "leann-core"
|
||||||
version = "0.3.4"
|
version = "0.3.3"
|
||||||
source = { editable = "packages/leann-core" }
|
source = { editable = "packages/leann-core" }
|
||||||
dependencies = [
|
dependencies = [
|
||||||
{ name = "accelerate" },
|
{ name = "accelerate" },
|
||||||
|
|||||||
Reference in New Issue
Block a user