* feat(embedding.py): add abstract base class Embedding and two concrete implementations SentenceTransformersEmbedding and OpenAIEmbedding

This commit is contained in:
NirantK
2023-07-11 19:32:53 +05:30
parent 964749d9f2
commit ea2584dd59

31
fastvector/embedding.py Normal file
View File

@@ -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