Upload vectordatabase.py

aimakerspace/vectordatabase.py

@@ -14,9 +14,9 @@ def cosine_similarity(vector_a: np.array, vector_b: np.array) -> float:
 
 
 class VectorDatabase:
-    def __init__(self, embedding_model: EmbeddingModel = None):
+    def __init__(self, embedding_model):
         self.vectors = defaultdict(np.array)
-        self.embedding_model = embedding_model or EmbeddingModel()
+        self.embedding_model = embedding_model
 
     def insert(self, key: str, vector: np.array) -> None:
         self.vectors[key] = vector
@@ -40,7 +40,8 @@ class VectorDatabase:
         distance_measure: Callable = cosine_similarity,
         return_as_text: bool = False,
     ) -> List[Tuple[str, float]]:
-        query_vector = self.embedding_model.get_embedding(query_text)
+        # query_vector = self.embedding_model.get_embedding(query_text)
+        query_vector = self.embedding_model.embed_query(query_text)
         results = self.search(query_vector, k, distance_measure)
         return [result[0] for result in results] if return_as_text else results
 
@@ -48,7 +49,8 @@ class VectorDatabase:
         return self.vectors.get(key, None)
 
     async def abuild_from_list(self, list_of_text: List[str]) -> "VectorDatabase":
-        embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
+        # embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
+        embeddings = await self.embedding_model.aembed_documents(list_of_text)
         for text, embedding in zip(list_of_text, embeddings):
             self.insert(text, np.array(embedding))
         return self