Embedding Improvement

1. move embedding function into llm_utils
2. add try feature with in embedding function
This commit is contained in:
Imccccc
2023-04-16 00:13:29 +08:00
parent 5e18bb4b61
commit f67b81e200
4 changed files with 44 additions and 10 deletions

View File

@@ -113,3 +113,35 @@ def create_chat_completion(
raise RuntimeError(f"Failed to get response after {num_retries} retries")
return response.choices[0].message["content"]
def create_embedding_with_ada(text) -> list:
"""Create a embedding with text-ada-002 using the OpenAI SDK"""
num_retries = 10
for attempt in range(num_retries):
backoff = 2 ** (attempt + 2)
try:
if CFG.use_azure:
return openai.Embedding.create(input=[text],
engine=CFG.get_azure_deployment_id_for_model("text-embedding-ada-002"),
)["data"][0]["embedding"]
else:
return openai.Embedding.create(input=[text], model="text-embedding-ada-002")[
"data"
][0]["embedding"]
except RateLimitError:
pass
except APIError as e:
if e.http_status == 502:
pass
else:
raise
if attempt == num_retries - 1:
raise
if CFG.debug_mode:
print(
Fore.RED + "Error: ",
f"API Bad gateway. Waiting {backoff} seconds..." + Fore.RESET,
)
time.sleep(backoff)