Upload folder using huggingface_hub

copy_chromadb.py

@@ -10,10 +10,11 @@ destination_collection_name = "vishnu_sahasranamam"
 source_collection = source_client.get_collection(source_collection_name)
 
 # Retrieve all data from the source collection
-source_data = source_collection.get(ids=source_collection.peek()["ids"]) # Efficiently get all IDs
+source_data = source_collection.get()
 
 # Create or get the destination collection
 if destination_client.get_or_create_collection(destination_collection_name):
+    print("Deleting existing collection", destination_collection_name)
     destination_client.delete_collection(destination_collection_name)
 
 destination_collection = destination_client.get_or_create_collection(
@@ -29,4 +30,6 @@ destination_collection.add(
     embeddings=source_data.get("embeddings") # Include embeddings if they exist in source
 )
 
-print("Collection copied successfully!")
+print("Collection copied successfully!")
+print("Total records in source collection = ", source_collection.count())
+print("Total records in destination collection = ", destination_collection.count())

db.py

@@ -28,8 +28,27 @@ class SanatanDatabase:
         )
 
     def search(self, collection_name: str, query: str, n_results=2):
+        logger.info("Searching for [%s] in [%s]", query, collection_name)
         collection = self.chroma_client.get_or_create_collection(name=collection_name)
         response = collection.query(
-            query_embeddings=[get_embedding(query)], n_results=n_results
+            query_embeddings=[get_embedding(query)], 
+            # query_texts=[query],
+            n_results=n_results
         )
         return response
+
+    def count(self, collection_name: str):
+        logger.info("Getting total records in [%s]", collection_name)
+        collection = self.chroma_client.get_or_create_collection(name=collection_name)
+        return collection.count()
+
+if __name__ == "__main__":
+    collection_name="vishnu_sahasranamam"
+    database = SanatanDatabase()
+    print("count = ", database.count(collection_name))
+    while True:
+        query = input("Search for: ")
+        if(query.strip() == ""):
+            break
+        response = database.search(collection_name=collection_name, query=query, n_results=1)
+        print(response["metadatas"][0][0]["translation"])