clean

.gitignore

@@ -1 +1,2 @@
-.env
+.env
+hf_to_chroma_ds

app.py

@@ -1,4 +1,8 @@
+import os
+from dotenv import load_dotenv
+
 import gradio as gr
+
 from langchain_chroma import Chroma
 from langchain.prompts import ChatPromptTemplate
 from langchain.chains import create_retrieval_chain, create_history_aware_retriever
@@ -6,25 +10,18 @@ from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import MessagesPlaceholder
 from langchain_community.chat_message_histories import ChatMessageHistory
 from langchain_core.runnables.history import RunnableWithMessageHistory
-import torch
-import chromadb
-from typing import List
 from langchain_core.documents import Document
 from langchain_core.retrievers import BaseRetriever
 from langchain_core.callbacks import CallbackManagerForRetrieverRun
 from langchain_core.vectorstores import VectorStoreRetriever
-
 from langchain_openai import ChatOpenAI
-from mixedbread_ai.client import MixedbreadAI
-
 from langchain.callbacks.tracers import ConsoleCallbackHandler
 from langchain_huggingface import HuggingFaceEmbeddings 
-import os
-# 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 dotenv import load_dotenv
+import chromadb
+from typing import List
+from mixedbread_ai.client import MixedbreadAI
 from tqdm import tqdm
 
 # Global params
@@ -37,18 +34,18 @@ MXBAI_API_KEY = os.environ.get("MXBAI_API_KEY")
 HF_TOKEN = os.environ.get("HF_TOKEN")
 HF_API_KEY = os.environ.get("HF_API_KEY")
 
-# Load the reranker model
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# MixedbreadAI Client
+# 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"  
 
 # 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)
+batched_ds = memoires_ds.batch(batch_size=50000)
 client = chromadb.Client()
-collection = client.get_or_create_collection(name="embeddings_mxbai") #, embedding_function=HuggingFaceEmbeddings(model_name=model_emb))
+collection = client.get_or_create_collection(name="embeddings_mxbai") 
 
-for batch in tqdm(batched_ds, desc="Processing dataset batches"): #, total=len(batched_ds)):
+for batch in tqdm(batched_ds, desc="Processing dataset batches"): 
     collection.add(
         ids=batch["id"],
         metadatas=batch["metadata"],