Merge pull request #1 from nikolas-rauscher/patch-1

Update OpenAI API Call freedmand#71

src/semantra/models.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 
 import numpy as np
-import openai
+from openai import OpenAI
 import tiktoken
 import torch
 from dotenv import load_dotenv
@@ -114,7 +114,7 @@ def __init__(
                 "OpenAI API key not set. Please set the OPENAI_API_KEY environment variable or create a `.env` file with the key in the current working directory or the Semantra directory, which is revealed by running `semantra --show-semantra-dir`."
             )
 
-        openai.api_key = os.getenv("OPENAI_API_KEY")
+        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
         self.model_name = model_name
         self.num_dimensions = num_dimensions
@@ -141,8 +141,8 @@ def get_text_chunks(self, _: str, tokens) -> "list[str]":
 
     def embed(self, tokens, offsets, _is_query=False) -> "list[list[float]]":
         texts = [tokens[i:j] for i, j in offsets]
-        response = openai.Embedding.create(model=self.model_name, input=texts)
-        return np.array([data["embedding"] for data in response["data"]])
+        response = self.client.embeddings.create(model=self.model_name, input=texts)
+        return np.array([embedding.embedding for embedding in response.data])
 
 
 def zero_if_none(x):