18 lines
585 B
Python
18 lines
585 B
Python
|
from config import Config
|
||
|
|
import cohere
|
||
|
|
import openai
|
||
|
|
|
||
|
|
class Embeddings:
|
||
|
|
def __init__(self):
|
||
|
|
self.cohere = cohere.Client(Config.COHERE_API_KEY)
|
||
|
|
|
||
|
|
def get_embedding(self, text: str, engine: str = "cohere"):
|
||
|
|
if engine == "cohere":
|
||
|
|
response = self.cohere.embed(texts=[text], model="small")
|
||
|
|
return response.embeddings[0]
|
||
|
|
else: # OpenAI fallback
|
||
|
|
response = openai.Embedding.create(
|
||
|
|
input=[text],
|
||
|
|
model="text-embedding-ada-002"
|
||
|
|
)
|
||
|
|
return response.data[0].embedding
|