collection

app.py

@@ -20,12 +20,12 @@ from mixedbread_ai.client import MixedbreadAI
 from langchain.callbacks.tracers import ConsoleCallbackHandler
 from langchain_huggingface import HuggingFaceEmbeddings 
 import os
-# from chroma_datasets.utils import import_into_chroma
-from hf_to_chroma_ds import import_into_chroma
+# from hf_to_chroma_ds import import_into_chroma
 from datasets import load_dataset
 from chromadb.utils import embedding_functions
-from hf_to_chroma_ds import Memoires_DS
+# from hf_to_chroma_ds import Memoires_DS
 from dotenv import load_dotenv
+from tqdm import tqdm
 
 # Global params
 CHROMA_PATH = "chromadb_mem10_mxbai_800_complete"
@@ -42,29 +42,20 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 mxbai_client = MixedbreadAI(api_key=MXBAI_API_KEY)
 model_emb = "mixedbread-ai/mxbai-embed-large-v1"  
 
-huggingface_ef = embedding_functions.huggingface_embedding_function.HuggingFaceEmbeddingFunction(
-    api_key=HF_API_KEY,
-    model_name=model_emb
-)
-
 # Set up ChromaDB
+memoires_ds = load_dataset("eliot-hub/memoires_vec_800", split="data", token=HF_TOKEN, streaming=True)
+batched_ds = memoires_ds.batch(batch_size=40000)
 client = chromadb.Client()
-# memoires_ds = load_dataset("eliot-hub/memoires_vec_800", split="data", token=HF_TOKEN)
-# client = chromadb.PersistentClient(path=os.path.join(os.path.abspath(os.getcwd()), "01_Notebooks", "RAG-ollama", "chatbot_actuariat_APP", CHROMA_PATH))
-
-# memoires_ds = Dataset(
-#     hf_data = None,
-#     hf_dataset_name = "eliot-hub/memoires_vec_800",
-#     embedding_function = huggingface_ef,
-#     embedding_function_instructions = None
-#     )
-
-
-collection = import_into_chroma(
-    chroma_client=client,
-    dataset=Memoires_DS,
-    embedding_function=huggingface_ef #Memoires_DS.embedding_function
+collection = client.get_or_create_collection(name="embeddings_mxbai") #, embedding_function=HuggingFaceEmbeddings(model_name=model_emb))
+
+for batch in tqdm(batched_ds, desc="Processing dataset batches"): #, total=len(batched_ds)):
+    collection.add(
+        ids=batch["id"],
+        metadatas=batch["metadata"],
+        documents=batch["document"],
+        embeddings=batch["embedding"],      
     )
+print(f"Collection complete: {collection.count()}")
 
 db = Chroma(
     client=client,
@@ -83,7 +74,7 @@ class Reranker(BaseRetriever):
         self, query: str, *, run_manager: CallbackManagerForRetrieverRun
     ) -> List[Document]:
         docs = self.retriever.invoke(query)
-        results = mxbai_client.reranking(model="mixedbread-ai/mxbai-rerank-large-v1", query=query, input=[doc.page_content for doc in docs], return_input=True, top_k=self.k)
+        results = mxbai_client.reranking(model=MODEL_RRK, query=query, input=[doc.page_content for doc in docs], return_input=True, top_k=self.k)
         return [Document(page_content=res.input) for res in results.data]
 
 # Set up reranker + LLM