Create core/embeddings.py

core/embeddings.py

@@ -0,0 +1,58 @@
+# core/embeddings.py
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from transformers import AutoTokenizer
+import faiss
+import numpy as np
+from typing import List, Dict
+import torch
+
+class DocumentEmbedder:
+    def __init__(self, model_name: str = "thenlper/gte-small"):
+        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
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.text_splitter = self._initialize_splitter()
+
+    def _initialize_splitter(self) -> RecursiveCharacterTextSplitter:
+        # Using markdown-optimized separators
+        MARKDOWN_SEPARATORS = [
+            "\n#{1,6} ",
+            "```\n",
+            "\n\\*\\*\\*+\n",
+            "\n---+\n",
+            "\n___+\n",
+            "\n\n",
+            "\n",
+            " ",
+            ""
+        ]
+        
+        return RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
+            self.tokenizer,
+            chunk_size=500,  # Adjusted for better semantic units
+            chunk_overlap=50,
+            add_start_index=True,
+            strip_whitespace=True,
+            separators=MARKDOWN_SEPARATORS
+        )
+
+    def process_documents(self, documents: List[Dict]) -> tuple:
+        """Process documents and return chunks and their embeddings."""
+        # Split documents into chunks
+        chunks = []
+        metadatas = []
+        
+        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])
+
+        # Generate embeddings
+        embeddings = self.embedding_model.embed_documents(chunks)
+        
+        return chunks, embeddings, metadatas