Update core/embeddings.py

core/embeddings.py

@@ -1,5 +1,5 @@
-# core/embeddings.py
-from langchain_community.embeddings import HuggingFaceEmbeddings
+# src/core/embeddings.py
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from transformers import AutoTokenizer
 import faiss
@@ -12,8 +12,13 @@ class DocumentEmbedder:
         self.model_name = model_name
         self.embedding_model = HuggingFaceEmbeddings(
             model_name=model_name,
-            model_kwargs={"device": "cuda" if torch.cuda.is_available() else "cpu"},
-            encode_kwargs={"normalize_embeddings": True}  # For cosine similarity
+            model_kwargs={
+                "device": "cuda" if torch.cuda.is_available() else "cpu"
+            },
+            encode_kwargs={
+                "normalize_embeddings": True,  # For cosine similarity
+                "batch_size": 32
+            }
         )
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.text_splitter = self._initialize_splitter()
@@ -34,7 +39,7 @@ class DocumentEmbedder:
         
         return RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
             self.tokenizer,
-            chunk_size=500,  # Adjusted for better semantic units
+            chunk_size=500,
             chunk_overlap=50,
             add_start_index=True,
             strip_whitespace=True,
@@ -50,7 +55,10 @@ class DocumentEmbedder:
         for doc in documents:
             doc_chunks = self.text_splitter.split_text(doc["content"])
             chunks.extend(doc_chunks)
-            metadatas.extend([{"source": doc["source"]} for _ in doc_chunks])
+            metadatas.extend([{
+                "source": doc["source"],
+                "chunk_index": i
+            } for i in range(len(doc_chunks))])
 
         # Generate embeddings
         embeddings = self.embedding_model.embed_documents(chunks)