File size: 11,609 Bytes
4467900 d7a2972 25fcabc d7a2972 86277c0 4467900 86277c0 b0174fa 4467900 d7a2972 4467900 d7a2972 b0174fa d7a2972 a8df5fb d7a2972 4467900 d7a2972 a8df5fb 86277c0 4467900 d7a2972 a8df5fb d7a2972 4467900 b0174fa d7a2972 b0174fa d7a2972 b0174fa d7a2972 25fcabc 86277c0 b0174fa d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 25fcabc d7a2972 b0174fa d7a2972 b0174fa d7a2972 b0174fa d7a2972 b0174fa d7a2972 b0174fa d7a2972 b0174fa d7a2972 a8df5fb d7a2972 b0174fa d7a2972 b0174fa d7a2972 b0174fa d7a2972 b0174fa |
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 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 |
import abc
import json
import os
from typing import Any, Generic, List, Optional, Tuple, TypeVar
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams
T = TypeVar("T", bound=dict[str, Any])
E = TypeVar("E")
class VectorStore(Generic[T, E], abc.ABC):
"""Abstract base class for a vector store.
A vector store is a key-value store that maps an ID to a vector embedding and a payload. The
payload can be any JSON-serializable object, e.g. a dictionary.
"""
INDEX_FILE = "vectors_index.json"
EMBEDDINGS_FILE = "vectors_data.npy"
PAYLOADS_FILE = "vectors_payloads.json"
@abc.abstractmethod
def _add(self, embedding: E, payload: T, emb_id: str) -> None:
"""Save an embedding with payload for a given ID."""
pass
@abc.abstractmethod
def _get(self, emb_id: str) -> Optional[E]:
"""Get the embedding for a given ID."""
pass
@abc.abstractmethod
def clear(self) -> None:
"""Clear the store."""
pass
def _get_emb_id(self, emb_id: Optional[str] = None, payload: Optional[T] = None) -> str:
if emb_id is None:
if payload is None:
raise ValueError("Either emb_id or payload must be provided.")
emb_id = json.dumps(payload, sort_keys=True)
return emb_id
def add(self, embedding: E, payload: T, emb_id: Optional[str] = None) -> None:
if emb_id is None:
emb_id = json.dumps(payload, sort_keys=True)
self._add(embedding=embedding, payload=payload, emb_id=emb_id)
def get(self, emb_id: Optional[str] = None, payload: Optional[T] = None) -> Optional[E]:
return self._get(emb_id=self._get_emb_id(emb_id=emb_id, payload=payload))
def has(self, emb_id: Optional[str] = None, payload: Optional[T] = None) -> bool:
return self.get(emb_id=emb_id, payload=payload) is not None
@abc.abstractmethod
def _retrieve_similar(
self, ref_id: str, top_k: Optional[int] = None, min_similarity: Optional[float] = None
) -> List[Tuple[T, float]]:
"""Retrieve IDs, payloads and the respective similarity scores with respect to the
reference entry. In the case that the reference entry is not in the store itself, an empty
list will be returned.
Args:
ref_id: The ID of the reference entry.
top_k: If provided, only the top-k most similar entries will be returned.
min_similarity: If provided, only entries with a similarity score greater or equal to
this value will be returned.
Returns:
A list of tuples consisting of the ID and the similarity score, sorted by similarity
score in descending order.
"""
pass
def retrieve_similar(
self, ref_id: Optional[str] = None, ref_payload: Optional[T] = None, **kwargs
) -> List[Tuple[T, float]]:
if not self.has(emb_id=ref_id, payload=ref_payload):
return []
return self._retrieve_similar(
ref_id=self._get_emb_id(emb_id=ref_id, payload=ref_payload), **kwargs
)
@abc.abstractmethod
def __len__(self):
pass
def _add_from_directory(self, directory: str) -> None:
with open(os.path.join(directory, self.INDEX_FILE), "r") as f:
index = json.load(f)
embeddings_np = np.load(os.path.join(directory, self.EMBEDDINGS_FILE))
with open(os.path.join(directory, self.PAYLOADS_FILE), "r") as f:
payloads = json.load(f)
for emb_id, emb, payload in zip(index, embeddings_np, payloads):
self._add(emb_id=emb_id, payload=payload, embedding=emb.tolist())
@abc.abstractmethod
def as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[T]]:
"""Return a tuple of indices, vectors and payloads."""
pass
def _save_to_directory(self, directory: str) -> None:
indices, vectors, payloads = self.as_indices_vectors_payloads()
np.save(os.path.join(directory, self.EMBEDDINGS_FILE), vectors)
with open(os.path.join(directory, self.PAYLOADS_FILE), "w") as f:
json.dump(payloads, f)
with open(os.path.join(directory, self.INDEX_FILE), "w") as f:
json.dump(indices, f)
def save_to_directory(self, directory: str) -> None:
"""Save the vector store to a directory."""
os.makedirs(directory, exist_ok=True)
self._save_to_directory(directory)
def load_from_directory(self, directory: str, replace: bool = False) -> None:
"""Load the vector store from a directory.
If `replace` is True, the current content of the store will be replaced.
"""
if replace:
self.clear()
self._add_from_directory(directory)
def vector_norm(vector: List[float]) -> float:
return sum(x**2 for x in vector) ** 0.5
def cosine_similarity(a: List[float], b: List[float]) -> float:
return sum(a * b for a, b in zip(a, b)) / (vector_norm(a) * vector_norm(b))
class SimpleVectorStore(VectorStore[T, List[float]]):
"""Simple in-memory vector store using a dictionary."""
def __init__(self):
self.vectors: dict[str, List[float]] = {}
self.payloads: dict[str, T] = {}
self._cache = {}
self._sim = cosine_similarity
def _add(self, embedding: E, payload: T, emb_id: str) -> None:
self.vectors[emb_id] = embedding
self.payloads[emb_id] = payload
def _get(self, emb_id: str) -> Optional[E]:
return self.vectors.get(emb_id)
def delete(self, emb_id: str) -> None:
if emb_id in self.vectors:
del self.vectors[emb_id]
del self.payloads[emb_id]
# remove from cache
self._cache = {k: v for k, v in self._cache.items() if emb_id not in k}
def clear(self) -> None:
self.vectors.clear()
self._cache.clear()
self.payloads.clear()
def __len__(self):
return len(self.vectors)
def _retrieve_similar(
self, ref_id: str, top_k: Optional[int] = None, min_similarity: Optional[float] = None
) -> List[Tuple[str, T, float]]:
ref_embedding = self.get(emb_id=ref_id)
if ref_embedding is None:
raise ValueError(f"Reference embedding '{ref_id}' not found.")
# calculate similarity to all embeddings
similarities = {}
for emb_id, embedding in self.vectors.items():
if (emb_id, ref_id) not in self._cache:
# use cosine similarity
self._cache[(emb_id, ref_id)] = self._sim(ref_embedding, embedding)
similarities[emb_id] = self._cache[(emb_id, ref_id)]
# sort by similarity
similar_entries = sorted(similarities.items(), key=lambda x: x[1], reverse=True)
if min_similarity is not None:
similar_entries = [
(emb_id, sim) for emb_id, sim in similar_entries if sim >= min_similarity
]
if top_k is not None:
similar_entries = similar_entries[:top_k]
return [(emb_id, self.payloads[emb_id], sim) for emb_id, sim in similar_entries]
def _save_to_directory(self, directory: str) -> None:
indices = list(self.vectors.keys())
with open(os.path.join(directory, self.INDEX_FILE), "w") as f:
json.dump(indices, f)
embeddings_np = np.array(list(self.vectors.values()))
np.save(os.path.join(directory, self.EMBEDDINGS_FILE), embeddings_np)
payloads = [self.payloads[idx] for idx in indices]
with open(os.path.join(directory, self.PAYLOADS_FILE), "w") as f:
json.dump(payloads, f)
def as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[T]]:
indices = list(self.vectors.keys())
embeddings_np = np.array(list(self.vectors.values()))
payloads = [self.payloads[idx] for idx in indices]
return indices, embeddings_np, payloads
class QdrantVectorStore(VectorStore[T, List[float]]):
"""Vector store using Qdrant as a backend."""
COLLECTION_NAME = "ADUs"
MAX_LIMIT = 100
def __init__(
self,
location: str = ":memory:",
vector_size: int = 768,
distance: Distance = Distance.COSINE,
):
self.client = QdrantClient(location=location)
self.emb_id2point_id = {}
self.point_id2emb_id = {}
self.client.create_collection(
collection_name=self.COLLECTION_NAME,
vectors_config=VectorParams(size=vector_size, distance=distance),
)
def __len__(self):
return self.client.get_collection(collection_name=self.COLLECTION_NAME).points_count
def _add(self, emb_id: str, payload: T, embedding: List[float]) -> None:
if emb_id in self.emb_id2point_id:
# update existing entry
point_id = self.emb_id2point_id[emb_id]
else:
# we use the length of the emb_id2point_id dict as the index,
# because we assume that, even when we delete an entry from
# the store, we do not delete it from the index
point_id = len(self.emb_id2point_id)
self.emb_id2point_id[emb_id] = point_id
self.point_id2emb_id[point_id] = emb_id
self.client.upsert(
collection_name=self.COLLECTION_NAME,
points=[PointStruct(id=point_id, vector=embedding, payload=payload)],
)
def _get(self, emb_id: str) -> Optional[List[float]]:
if emb_id not in self.emb_id2point_id:
return None
points = self.client.retrieve(
collection_name=self.COLLECTION_NAME,
ids=[self.emb_id2point_id[emb_id]],
with_vectors=True,
)
if len(points) == 0:
return None
elif len(points) == 1:
return points[0].vector
else:
raise ValueError(f"Multiple points found for ID '{emb_id}'.")
def _retrieve_similar(
self, ref_id: str, top_k: Optional[int] = None, min_similarity: Optional[float] = None
) -> List[Tuple[str, T, float]]:
similar_entries = self.client.recommend(
collection_name=self.COLLECTION_NAME,
positive=[self.emb_id2point_id[ref_id]],
limit=top_k or self.MAX_LIMIT,
score_threshold=min_similarity,
)
return [
(self.point_id2emb_id[entry.id], entry.payload, entry.score)
for entry in similar_entries
]
def clear(self) -> None:
vectors_config = self.client.get_collection(
collection_name=self.COLLECTION_NAME
).vectors_config
self.client.delete_collection(collection_name=self.COLLECTION_NAME)
self.client.create_collection(
collection_name=self.COLLECTION_NAME,
vectors_config=vectors_config,
)
self.emb_id2point_id.clear()
self.point_id2emb_id.clear()
def as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[T]]:
num_entries = self.client.get_collection(collection_name=self.COLLECTION_NAME).points_count
data, point_ids = self.client.scroll(
collection_name=self.COLLECTION_NAME, with_vectors=True, limit=num_entries
)
vectors_np = np.array([point.vector for point in data])
payloads = [point.payload for point in data]
emb_ids = [self.point_id2emb_id[point.id] for point in data]
return emb_ids, vectors_np, payloads
|