changed embedding model to finetuned model

src/app.py

@@ -18,7 +18,7 @@ from src.core.parsing import read_file
 VECTOR_STORE = "faiss"
 MODEL = "openai"
 EMBEDDING = "openai"
-MODEL = "gpt-3.5-turbo-16k"
+MODEL = "gpt-4"
 K = 5
 USE_VERBOSE = True
 API_KEY = os.environ["OPENAI_API_KEY"]
@@ -52,12 +52,7 @@ def getretriever():
       print(e)
 
   chunked_file = chunk_file(file, chunk_size=512, chunk_overlap=0)
-  folder_index = embed_files(
-    files=[chunked_file],
-    embedding=EMBEDDING,
-    vector_store=VECTOR_STORE,
-    openai_api_key=API_KEY,
-  )
+  folder_index = embed_files(files=[chunked_file])
   return folder_index.index.as_retriever(verbose=True, search_type="similarity", search_kwargs={"k": K})
 
 retriever = getretriever()

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 HuggingFaceEmbeddings
 from langchain.embeddings.base import Embeddings
 from langchain.vectorstores import VectorStore
 from langchain.vectorstores.faiss import FAISS
@@ -47,30 +47,18 @@ class FolderIndex:
         return cls(files=files, index=index)
 
 
-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.")
-
+def embed_files(files: List[File]) -> FolderIndex:
+    model_name = "adriancowham/letstalk-embed-gte-small"
+    model_kwargs = {'device': 'cpu'}
+    encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
+    print("Loading model...")
+    model_norm = HuggingFaceEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs
+    )
+    print("Model loaded.")
+    embeddings = model_norm
     return FolderIndex.from_files(
-        files=files, embeddings=_embeddings, vector_store=_vector_store
+        files=files, embeddings=embeddings, vector_store=FAISS
     )