Spaces:
Runtime error
Runtime error
File size: 3,202 Bytes
fca155a |
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 |
import numpy as np
from pathlib import Path
from typing import List, Tuple, Dict, Any, Optional
from sklearn.metrics.pairwise import cosine_similarity
import pickle
class VectorIndex:
"""
In-memory Vector Database.
This acts as the 'Long Term Memory' for visual concepts.
It maps a Timestamp (when something happened) to a Vector (what it looked like).
"""
def __init__(self, index_file_path: Path):
self.file_path = index_file_path
self.timestamps: List[float] = []
self.embedding_matrix: Optional[np.ndarray] = None
self.metadata_store: List[Dict[str, Any]] = []
# Load existing index if available
if self.file_path.exists():
self.load()
def add(self, timestamp_seconds: float, vector: np.ndarray, extra_data: Dict[str, Any] = None):
"""Adds a new memory entry (timestamp + vector)."""
self.timestamps.append(timestamp_seconds)
self.metadata_store.append(extra_data or {})
# Normalize the vector to length 1.
# This is crucial so that 'Cosine Similarity' is just a Dot Product (faster).
vector_norm = np.linalg.norm(vector)
if vector_norm > 0:
vector = vector / vector_norm
if self.embedding_matrix is None:
self.embedding_matrix = vector.reshape(1, -1)
else:
self.embedding_matrix = np.vstack([self.embedding_matrix, vector])
def search(self, query_vector: np.ndarray, top_k: int = 5) -> List[Tuple[float, float]]:
"""
Finds the moments in the video that are most similar to the query.
Returns:
A list of tuples: (timestamp_seconds, similarity_score)
"""
if self.embedding_matrix is None:
return []
# Normalize the query too
query_norm = np.linalg.norm(query_vector)
if query_norm > 0:
query_vector = query_vector / query_norm
# Calculate similarity against ALL stored memories at once
similarity_scores = cosine_similarity(query_vector.reshape(1, -1), self.embedding_matrix)[0]
# Sort by highest score first
best_indices = np.argsort(similarity_scores)[::-1][:top_k]
results = []
for index in best_indices:
score = float(similarity_scores[index])
time_point = self.timestamps[index]
results.append((time_point, score))
return results
def save(self):
"""Persists the index to the disk using Pickle."""
data_packet = {
"timestamps": self.timestamps,
"vectors": self.embedding_matrix,
"metadata": self.metadata_store
}
with open(self.file_path, "wb") as f:
pickle.dump(data_packet, f)
def load(self):
"""Loads the index from disk."""
with open(self.file_path, "rb") as f:
data_packet = pickle.load(f)
self.timestamps = data_packet["timestamps"]
self.embedding_matrix = data_packet["vectors"]
self.metadata_store = data_packet.get("metadata", []) |