jim/morphik-core
1
0
Fork 0
mirror of https://github.com/james-m-jordan/morphik-core.git synced 2025-05-09 19:32:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

Speed up multivector store

Browse Source 
Co-authored-by: Arnav Agrawal <aa779@cornell.edu>
This commit is contained in:
Adityavardhan Agrawal 2025-04-20 20:02:42 -07:00
parent 89633ce761
commit bce0e1cfe1
1 changed files with 16 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
core/vector_store/multi_vector_store.py
View File

@ -181,7 +181,8 @@ class MultiVectorStore(BaseVectorStore):
""" """
CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) RETURNS double precision AS $$ CREATE OR REPLACE FUNCTION max_sim(document bit[], query bit[]) RETURNS double precision AS $$
WITH queries AS ( WITH queries AS (
SELECT row_number() OVER () AS query_number, * FROM (SELECT unnest(query) AS query) AS foo SELECT row_number() OVER () AS query_number, *
FROM (SELECT unnest(query) AS query) AS foo
), ),
documents AS ( documents AS (
SELECT unnest(document) AS document SELECT unnest(document) AS document
@ -189,7 +190,8 @@ class MultiVectorStore(BaseVectorStore):
similarities AS ( similarities AS (
SELECT SELECT
query_number, query_number,
1.0 - (bit_count(document # query)::float / greatest(bit_length(query), 1)::float) AS similarity 1.0 - (bit_count(document # query)::float /
greatest(bit_length(query), 1)::float) AS similarity
FROM queries CROSS JOIN documents FROM queries CROSS JOIN documents
), ),
max_similarities AS ( max_similarities AS (
@ -217,10 +219,6 @@ class MultiVectorStore(BaseVectorStore):
if isinstance(embeddings, list) and not isinstance(embeddings[0], np.ndarray): if isinstance(embeddings, list) and not isinstance(embeddings[0], np.ndarray):
embeddings = np.array(embeddings) embeddings = np.array(embeddings)
# Add this check to ensure pgvector is registered for the connection
with self.get_connection() as conn:
register_vector(conn)
return [Bit(embedding > 0) for embedding in embeddings] return [Bit(embedding > 0) for embedding in embeddings]
async def store_embeddings(self, chunks: List[DocumentChunk]) -> Tuple[bool, List[str]]: async def store_embeddings(self, chunks: List[DocumentChunk]) -> Tuple[bool, List[str]]:
@ -230,7 +228,7 @@ class MultiVectorStore(BaseVectorStore):
return True, [] return True, []
stored_ids = [] stored_ids = []
with self.get_connection() as conn:
for chunk in chunks: for chunk in chunks:
# Ensure embeddings exist # Ensure embeddings exist
if not hasattr(chunk, "embedding") or chunk.embedding is None: if not hasattr(chunk, "embedding") or chunk.embedding is None:
@ -244,7 +242,7 @@ class MultiVectorStore(BaseVectorStore):
binary_embeddings = self._binary_quantize(embeddings) binary_embeddings = self._binary_quantize(embeddings)
# Insert into database with retry logic # Insert into database with retry logic
with self.get_connection() as conn:
conn.execute( conn.execute(
""" """
INSERT INTO multi_vector_embeddings INSERT INTO multi_vector_embeddings