Spaces:
Sleeping
Sleeping
File size: 3,486 Bytes
b16a546 | 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 | """
vector_db.py — A minimal vector database backed by a JSON file.
Flow:
1. Embed text using a sentence-transformer model (runs locally).
2. Store the embedding + metadata as a record in data.json.
3. At query time, embed the query and rank all records by cosine similarity.
"""
import json
import math
import os
from typing import Any
class VectorDB:
def __init__(self, db_path: str = "data.json", model_name: str = "all-MiniLM-L6-v2"):
self.db_path = db_path
self._model_name = model_name
self._model = None # loaded lazily so imports stay fast
self._records: list[dict] = []
if os.path.exists(db_path):
self.load()
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _get_model(self):
if self._model is None:
from sentence_transformers import SentenceTransformer
print(f"[VectorDB] Loading model '{self._model_name}' (first run downloads ~80 MB)...")
self._model = SentenceTransformer(self._model_name)
print("[VectorDB] Model ready.\n")
return self._model
@staticmethod
def _cosine_similarity(a: list[float], b: list[float]) -> float:
dot = sum(x * y for x, y in zip(a, b))
mag_a = math.sqrt(sum(x * x for x in a))
mag_b = math.sqrt(sum(x * x for x in b))
if mag_a == 0 or mag_b == 0:
return 0.0
return dot / (mag_a * mag_b)
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def add(self, text: str, metadata: dict[str, Any] | None = None) -> None:
"""Embed `text` and append a record to the in-memory database."""
model = self._get_model()
embedding: list[float] = model.encode(text).tolist()
record = {
"id": len(self._records),
"text": text,
"metadata": metadata or {},
"embedding": embedding,
}
self._records.append(record)
def save(self) -> None:
"""Persist all records (including raw embeddings) to JSON."""
with open(self.db_path, "w", encoding="utf-8") as f:
json.dump(self._records, f, indent=2)
print(f"[VectorDB] Saved {len(self._records)} records -> {self.db_path}")
def load(self) -> None:
"""Load records from the JSON file (embeddings included)."""
with open(self.db_path, encoding="utf-8") as f:
self._records = json.load(f)
print(f"[VectorDB] Loaded {len(self._records)} records <- {self.db_path}")
def search(self, query: str, top_k: int = 5) -> list[dict]:
"""
Embed the query and return the top_k most similar records,
ordered by cosine similarity (highest first).
"""
if not self._records:
return []
model = self._get_model()
query_vec: list[float] = model.encode(query).tolist()
scored = [
{
"score": self._cosine_similarity(query_vec, rec["embedding"]),
"id": rec["id"],
"text": rec["text"],
"metadata": rec["metadata"],
}
for rec in self._records
]
scored.sort(key=lambda r: r["score"], reverse=True)
return scored[:top_k]
|