Merge pull request exowanderer#5 from philippesaade-wmde/updated_merge_conflicts

- Improved cache db operations and fixed errors involve db insertion
- Improved data storage by storing the smaller base64 data per embedding, which results directly from JinaAI's API
- Refactored handling of DataStax and JinaAI errors
- Fixed duplicate ID errors
- updated our caching and merging of cache
- created migrate_db routine to update from original to smaller db

Author: exowanderer

docker/7_Create_Prototype/run.py

@@ -88,12 +88,19 @@ def process_items(queue, progress_bar):
             break  # Exit condition for worker processes
 
         item_id = item['id']
-        item_label = textifier.get_label(item_id, json.loads(item['labels']))
+
+        item_label = textifier.get_label(
+            item_id,
+            json.loads(item['labels'])
+        )
+
         item_description = textifier.get_description(
             item_id,
             json.loads(item['descriptions'])
         )
-        item_aliases = textifier.get_aliases(json.loads(item['aliases']))
+        item_aliases = textifier.get_aliases(
+            json.loads(item['aliases'])
+        )
 
         if item_label is not None:
             # TODO: Verify: If label does not exist, then skip item
@@ -128,16 +135,15 @@ def process_items(queue, progress_bar):
 
                 graph_store.add_document(
                     id=f"{item_id}_{LANGUAGE}_{chunk_i+1}",
+
                     text=chunk,
                     metadata=metadata
                 )
 
         progress_bar.value += 1
 
-    while True:
-        # Leftover Maintenance: Ensure that the batch is emptied out
-        if not graph_store.push_batch():  # Stop when batch is empty
-            break
+
+    graph_store.push_all()
 
 
 if __name__ == "__main__":

src/merge_cache.py

@@ -0,0 +1,91 @@
+import sqlite3
+import glob
+import os
+import base64
+from tqdm import tqdm
+
+# Define database file pattern
+db_files = glob.glob("../data/Wikidata/sqlite_cacheembeddings_*.db")
+
+# Define the target merged database
+merged_db = "../data/Wikidata/sqlite_cacheembeddings_merged.db"
+TABLE_NAME = "wikidata_prototype"
+
+# Batch size for processing
+BATCH_SIZE = 1000  # Adjust based on performance needs
+
+# Create the merged database connection
+conn_merged = sqlite3.connect(merged_db)
+cursor_merged = conn_merged.cursor()
+
+# Create table in the merged database if it doesn't exist
+cursor_merged.execute(f"""
+CREATE TABLE IF NOT EXISTS {TABLE_NAME} (
+    id TEXT PRIMARY KEY,
+    embedding TEXT
+);
+""")
+conn_merged.commit()
+
+# Helper function to check if a string is a valid Base64 encoding
+def is_valid_base64(s):
+    try:
+        if not s or not isinstance(s, str):
+            return False
+        base64.b64decode(s, validate=True)
+        return True
+    except Exception:
+        return False
+
+# Loop through all source databases
+for db_file in db_files:
+    print(f"Processing {db_file}...")
+
+    # Connect to the current database
+    conn_src = sqlite3.connect(db_file)
+    cursor_src = conn_src.cursor()
+
+    # Get total record count for progress tracking
+    cursor_src.execute(f"SELECT COUNT(*) FROM {TABLE_NAME}")
+    total_records = cursor_src.fetchone()[0]
+
+    # Fetch records in batches
+    offset = 0
+    with tqdm(total=total_records,
+              desc=f"Merging {db_file}", unit="records") as pbar:
+        while True:
+            cursor_src.execute(f"SELECT id, embedding FROM {TABLE_NAME} LIMIT {BATCH_SIZE} OFFSET {offset}")
+            records = cursor_src.fetchall()
+            if not records:
+                break  # No more records to process
+
+            # Prepare batch for insertion
+            batch_data = []
+            for id_, embedding in records:
+                if embedding and embedding.strip() and is_valid_base64(embedding):
+                    cursor_merged.execute(f"SELECT embedding FROM {TABLE_NAME} WHERE id = ?", (id_,))
+                    existing = cursor_merged.fetchone()
+
+                    if existing is None or not is_valid_base64(existing[0]):
+                        batch_data.append((id_, embedding, embedding))  # Prepare for bulk insert
+
+            # Perform batch insert/update
+            if batch_data:
+                cursor_merged.executemany(
+                    f"""
+                    INSERT INTO {TABLE_NAME} (id, embedding)
+                    VALUES (?, ?)
+                    ON CONFLICT(id) DO UPDATE SET embedding = ?
+                    """,
+                    batch_data
+                )
+                conn_merged.commit()
+
+            offset += BATCH_SIZE  # Move to the next batch
+            pbar.update(len(records))  # Update tqdm progress bar
+
+    conn_src.close()
+
+# Close merged database connection
+conn_merged.close()
+print(f"Merge completed! Combined database saved as {merged_db}")

src/migrate_cache.py

@@ -0,0 +1,77 @@
+import sqlite3
+import json
+import base64
+import numpy as np
+from tqdm import tqdm
+
+DB_PATH = "../data/Wikidata/sqlite_cacheembeddings.db"
+TABLE_NAME = "wikidata_prototype"  # Change this to match your actual table name
+BATCH_SIZE = 5000  # Process in smaller batches to avoid memory overload
+
+def convert_embeddings():
+    """
+    Convert JSON-stored embeddings into Base64-encoded binary format in batches.
+    Uses `fetchmany(BATCH_SIZE)` to process records iteratively.
+    """
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+
+    # Check if the embedding column exists (sanity check)
+    cursor.execute(f"PRAGMA table_info({TABLE_NAME})")
+    columns = [row[1] for row in cursor.fetchall()]
+    if "embedding" not in columns:
+        print("Error: 'embedding' column does not exist in the table!")
+        return
+
+    # Count total records for progress tracking
+    cursor.execute(f"SELECT COUNT(*) FROM {TABLE_NAME}")
+    total_records = cursor.fetchone()[0]
+
+    print(f"Total records to process: {total_records}")
+
+    # Fetch records in batches using an iterator
+    offset = 0
+    with tqdm(total=total_records, desc="Converting embeddings", unit="record") as pbar:
+        while True:
+            cursor.execute(f"SELECT id, embedding FROM {TABLE_NAME} LIMIT {BATCH_SIZE} OFFSET {offset}")
+            records = cursor.fetchall()
+            if not records:
+                break  # Stop when there are no more records
+
+            updated_records = []
+            for id, json_embedding in records:
+                if json_embedding:
+                    try:
+                        # Convert JSON string to list of floats
+                        embedding_list = json.loads(json_embedding)
+
+                        # Convert list of floats to Base64-encoded binary
+                        binary_data = np.array(embedding_list, dtype=np.float32).tobytes()
+                        base64_embedding = base64.b64encode(binary_data).decode('utf-8')
+
+                        updated_records.append((base64_embedding, id))
+                    except Exception as e:
+                        pass
+
+                pbar.update(1)  # Update progress bar for each record processed
+
+            # Update database in batches
+            if updated_records:
+                cursor.executemany(
+                    f"UPDATE {TABLE_NAME} SET embedding = ? WHERE id = ?",
+                    updated_records
+                )
+                conn.commit()  # Commit every batch
+
+            offset += BATCH_SIZE  # Move to next batch
+
+    print("Optimizing database with VACUUM...")
+    cursor.execute("VACUUM;")
+    conn.commit()
+
+    print("Migration completed successfully.")
+
+    conn.close()
+
+if __name__ == "__main__":
+    convert_embeddings()

src/wikidataCache.py

@@ -5,6 +5,8 @@
 
 import os
 import json
+import base64
+import numpy as np
 
 """
 SQLite database setup for caching the query embeddings for a faster
@@ -36,19 +38,28 @@
 Base = declarative_base()
 Session = sessionmaker(bind=engine)
 
+class EmbeddingType(TypeDecorator):
+    """Custom SQLAlchemy type for storing embeddings as Base64 strings in SQLite."""
 
-class JSONType(TypeDecorator):
-    """Custom SQLAlchemy type for JSON storage in SQLite."""
     impl = Text
 
     def process_bind_param(self, value, dialect):
-        if value is not None:
-            return json.dumps(value, separators=(',', ':'))
+        """Convert a list of floats (embedding) to a Base64 string before storing."""
+        if value is not None and isinstance(value, list):
+            # Convert list to binary
+            binary_data = np.array(value, dtype=np.float32).tobytes()
+            # Encode to Base64 string
+            return base64.b64encode(binary_data).decode('utf-8')
         return None
 
     def process_result_value(self, value, dialect):
+        """Convert a Base64 string back to a list of floats when retrieving."""
         if value is not None:
-            return json.loads(value)
+            # Decode Base64
+            binary_data = base64.b64decode(value)
+            # Convert back to float32 list
+            embedding_array = np.frombuffer(binary_data, dtype=np.float32)
+            return embedding_array.tolist()
         return None
 
 
@@ -59,7 +70,7 @@ class CacheEmbeddings(Base):
         __tablename__ = table_name
 
         id = Column(Text, primary_key=True)
-        embedding = Column(JSONType)
+        embedding = Column(EmbeddingType)
 
         @staticmethod
         def add_cache(id, embedding):
@@ -99,6 +110,13 @@ def add_bulk_cache(data):
             - bool: True if the operation was successful, False otherwise.
             """
             worked = False
+            embeddingtype = EmbeddingType()
+            for i in range(len(data)):
+                data[i]['embedding'] = embeddingtype.process_bind_param(
+                    data[i]['embedding'],
+                    None
+                )
+
             with Session() as session:
                 exec_text = text(
                     f"""

src/wikidataRetriever.py

@@ -17,6 +17,7 @@ def __init__(self, datastax_token, collection_name, model='jina', batch_size=8,
         from langchain_astradb import AstraDBVectorStore
         from astrapy.info import CollectionVectorServiceOptions
         from astrapy import DataAPIClient
+        from astrapy.exceptions import InsertManyException
         from multiprocessing import Queue
 
         from transformers import AutoTokenizer
@@ -30,6 +31,7 @@ def __init__(self, datastax_token, collection_name, model='jina', batch_size=8,
         self.model = model
         self.collection_name = collection_name
         self.doc_batch = Queue()
+        self.InsertManyException = InsertManyException
 
         self.cache_on = (cache_embeddings is not None)
         if self.cache_on:
@@ -126,23 +128,29 @@ def push_batch(self):
         if len(docs) == 0:
             return False
 
-        try:
-            vectors = self.embeddings.embed_documents(
-                [doc['content'] for doc in docs]
-            )
-            self.graph_store.insert_many(docs, vectors=vectors)
-        except Exception as e:
-            print(e)
-
-            # Put the documents back in the Queue and try again later.
-            for doc in docs:
-                self.doc_batch.put(doc)
+        while True:
+            try:
+                vectors = self.embeddings.embed_documents(
+                    [doc['content'] for doc in docs]
+                )
+                break
+            except Exception as e:
+                print(e)
+                time.sleep(3)
 
-            return False
+        while True:
+            try:
+                self.graph_store.insert_many(docs, vectors=vectors)
+                break
+            except self.InsertManyException as e:
+                break
+            except Exception as e:
+                print(e)
+                time.sleep(3)
 
         self.cache_model.add_bulk_cache([{
             'id': docs[i]['_id'],
-            'embedding': json.dumps(vectors[i], separators=(',', ':'))}
+            'embedding': vectors[i]}
             for i in range(len(docs))])
 
         return True