From ea2584dd59510fbcbb335de0719368f47c5e96f0 Mon Sep 17 00:00:00 2001
From: NirantK <nirant.bits@gmail.com>
Date: Tue, 11 Jul 2023 19:32:53 +0530
Subject: [PATCH] * feat(embedding.py): add abstract base class Embedding and
 two concrete implementations SentenceTransformersEmbedding and
 OpenAIEmbedding

---
 fastvector/embedding.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)
 create mode 100644 fastvector/embedding.py

diff --git a/fastvector/embedding.py b/fastvector/embedding.py
new file mode 100644
index 0000000..057c975
--- /dev/null
+++ b/fastvector/embedding.py
@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+
+
+class Embedding(ABC):
+    @abstractmethod
+    def encode(self, texts):
+        pass
+
+
+class SentenceTransformersEmbedding(Embedding):
+    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError:
+            raise ImportError("Please install the sentence-transformers package to use this method.")
+        self.model = SentenceTransformer(model_name)
+
+    def encode(self, texts):
+        return self.model.encode(texts)
+
+
+class OpenAIEmbedding(Embedding):
+    def __init__(self):
+        # Initialize your OpenAI model here
+        # self.model = ...
+        ...
+
+    def encode(self, texts):
+        # Use your OpenAI model to encode the texts
+        # return self.model.encode(texts)
+        raise NotImplementedError