changed embeddings

src/core/embedding.py

@@ -1,7 +1,7 @@
 from typing import List, Type
 
 from langchain.docstore.document import Document
-from langchain.embeddings import OpenAIEmbeddings
+from langchain.embeddings import HuggingFaceBgeEmbeddings
 from langchain.embeddings.base import Embeddings
 from langchain.vectorstores import VectorStore
 from langchain.vectorstores.faiss import FAISS
@@ -50,27 +50,16 @@ class FolderIndex:
 def embed_files(
     files: List[File], embedding: str, vector_store: str, **kwargs
 ) -> FolderIndex:
-    """Embeds a collection of files and stores them in a FolderIndex."""
-
-    supported_embeddings: dict[str, Type[Embeddings]] = {
-        "openai": OpenAIEmbeddings,
-        "debug": FakeEmbeddings,
-    }
-    supported_vector_stores: dict[str, Type[VectorStore]] = {
-        "faiss": FAISS,
-        "debug": FakeVectorStore,
-    }
-
-    if embedding in supported_embeddings:
-        _embeddings = supported_embeddings[embedding](**kwargs)
-    else:
-        raise NotImplementedError(f"Embedding {embedding} not supported.")
-
-    if vector_store in supported_vector_stores:
-        _vector_store = supported_vector_stores[vector_store]
-    else:
-        raise NotImplementedError(f"Vector store {vector_store} not supported.")
-
+    model_name = "BAAI/bge-small-en"
+    model_kwargs = {'device': 'mps'}
+    encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
+    model_norm = HuggingFaceBgeEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs
+    )
+    # embeddings = OpenAIEmbeddings
+    embeddings = model_norm
     return FolderIndex.from_files(
-        files=files, embeddings=_embeddings, vector_store=_vector_store
+        files=files, embeddings=embeddings, vector_store=FAISS
     )