From 09b708566b32c11ec0cbbf64501deb886751902b Mon Sep 17 00:00:00 2001
From: Nikolas <75223225+nikolas-rauscher@users.noreply.github.com>
Date: Fri, 1 Mar 2024 21:57:31 +0100
Subject: [PATCH] Update OpenAI API Call #71

Fix #71
---
 src/semantra/models.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/semantra/models.py b/src/semantra/models.py
index 5ff77b1..36d0080 100644
--- a/src/semantra/models.py
+++ b/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):