aimakerspace/openai_utils/embedding.py

import os
import openai
import numpy as np
from typing import List, Union
import asyncio

class EmbeddingModel:
    def __init__(self, model_name: str = "text-embedding-3-small"):
        self.model_name = model_name
        self.openai_api_key = os.getenv("OPENAI_API_KEY")
        if self.openai_api_key is None:
            raise ValueError("OPENAI_API_KEY is not set")

    def get_embedding(self, text: str) -> np.ndarray:
        openai.api_key = self.openai_api_key
        response = openai.Embedding.create(
            model=self.model_name,
            input=text
        )
        return np.array(response.data[0].embedding)

    async def async_get_embeddings(self, list_of_text: List[str]) -> List[np.ndarray]:
        openai.api_key = self.openai_api_key
        response = await openai.Embedding.acreate(
            model=self.model_name,
            input=list_of_text
        )
        return [np.array(item.embedding) for item in response.data]