Update app/backend/chromadb_utils.py

app/backend/chromadb_utils.py

@@ -1,8 +1,49 @@
 import chromadb
 from chromadb.utils import embedding_functions
 from .text_processing import text_chunking
+import torch
+from transformers import AutoTokenizer, AutoModel
 
 
+class SimpleEmbeddingFunction:
+    def __init__(self, model_path):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
+        self.model = AutoModel.from_pretrained(model_path, local_files_only=True)
+        self.model.eval()
+    
+    def __call__(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        inputs = self.tokenizer(
+            texts, 
+            padding=True, 
+            truncation=True, 
+            return_tensors='pt',
+            max_length=384
+        )
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            embeddings = outputs.last_hidden_state.mean(dim=1)
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        
+        return embeddings.numpy().tolist()
+
+
+def initialize_chromadb(model_name, model_path):
+    import chromadb
+    
+    if not os.path.exists(model_path):
+        raise Exception(f"Local model not found at {model_path}")
+    
+    embedding_func = SimpleEmbeddingFunction(model_path)
+    client = chromadb.Client()
+    
+    return client, embedding_func
+
+
+"""
 def initialize_chromadb(EMBEDDING_MODEL, local_model_path=None):
     """
     Initialize ChromaDB client and embedding function, using a local model path if provided.
@@ -19,6 +60,7 @@ def initialize_chromadb(EMBEDDING_MODEL, local_model_path=None):
         )
 
     return client, embedding_func
+"""
 
 
 def initialize_collection(client, embedding_func, collection_name):