File size: 4,955 Bytes
d61821a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """Pinned dense-index backends used by E01 and E05."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import sqlite3
import time
from typing import Sequence
import faiss
import numpy as np
import sqlite_vec
from .components import Candidate
from .retrieval import DenseRetriever, _candidate
@dataclass(frozen=True, slots=True)
class VectorBackendStats:
backend: str
build_seconds: float
index_disk_bytes: int
parameters: dict[str, int | str]
def _matrix(dense: DenseRetriever) -> np.ndarray:
return np.asarray([list(item) for item in dense.vectors], dtype="float32")
class FaissFlatRetriever:
def __init__(self, dense: DenseRetriever):
started = time.monotonic()
self.dense = dense
matrix = _matrix(dense)
self.index = faiss.IndexFlatIP(dense.spec.vector_dimension)
self.index.add(matrix)
self.stats = VectorBackendStats(
backend="faiss_index_flat_ip",
build_seconds=time.monotonic() - started,
index_disk_bytes=0,
parameters={"dimension": dense.spec.vector_dimension, "metric": "inner_product"},
)
def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
vector = np.asarray([list(self.dense.query_vector(query))], dtype="float32")
scores, indices = self.index.search(vector, min(limit, len(self.dense.chunks)))
return tuple(
_candidate(self.dense.chunks[int(index)], "dense_faiss_flat", float(score))
for score, index in zip(scores[0], indices[0])
if index >= 0
)
class FaissHNSWRetriever:
def __init__(
self,
dense: DenseRetriever,
neighbors: int = 32,
ef_construction: int = 80,
ef_search: int = 64,
):
started = time.monotonic()
self.dense = dense
matrix = _matrix(dense)
self.index = faiss.IndexHNSWFlat(
dense.spec.vector_dimension,
neighbors,
faiss.METRIC_INNER_PRODUCT,
)
self.index.hnsw.efConstruction = ef_construction
self.index.hnsw.efSearch = ef_search
self.index.add(matrix)
self.stats = VectorBackendStats(
backend="faiss_hnsw_ip",
build_seconds=time.monotonic() - started,
index_disk_bytes=0,
parameters={
"dimension": dense.spec.vector_dimension,
"metric": "inner_product",
"neighbors": neighbors,
"ef_construction": ef_construction,
"ef_search": ef_search,
},
)
def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
vector = np.asarray([list(self.dense.query_vector(query))], dtype="float32")
scores, indices = self.index.search(vector, min(limit, len(self.dense.chunks)))
return tuple(
_candidate(self.dense.chunks[int(index)], "dense_faiss_hnsw", float(score))
for score, index in zip(scores[0], indices[0])
if index >= 0
)
class SQLiteVecRetriever:
def __init__(self, dense: DenseRetriever, path: Path):
started = time.monotonic()
self.dense = dense
self.path = path
path.parent.mkdir(parents=True, exist_ok=True)
self.connection = sqlite3.connect(path)
self.connection.enable_load_extension(True)
sqlite_vec.load(self.connection)
self.connection.enable_load_extension(False)
self.connection.execute("DROP TABLE IF EXISTS vec_items")
self.connection.execute(
f"CREATE VIRTUAL TABLE vec_items USING vec0("
f"embedding float[{dense.spec.vector_dimension}] distance_metric=cosine)"
)
self.connection.executemany(
"INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)",
((index + 1, vector.tobytes()) for index, vector in enumerate(dense.vectors)),
)
self.connection.commit()
self.stats = VectorBackendStats(
backend="sqlite_vec_exact",
build_seconds=time.monotonic() - started,
index_disk_bytes=path.stat().st_size,
parameters={"dimension": dense.spec.vector_dimension, "metric": "cosine"},
)
def retrieve(self, query: str, limit: int) -> Sequence[Candidate]:
vector = self.dense.query_vector(query)
rows = self.connection.execute(
"SELECT rowid, distance FROM vec_items WHERE embedding MATCH ? AND k = ? ORDER BY distance",
(vector.tobytes(), min(limit, len(self.dense.chunks))),
).fetchall()
return tuple(
_candidate(
self.dense.chunks[int(rowid) - 1],
"dense_sqlite_vec",
1.0 - float(distance),
)
for rowid, distance in rows
)
def close(self) -> None:
self.connection.close()
|