如何修复从 Databricks 对 OpenAI 嵌入的异步调用？

Question

我正在尝试使用 OpenAI Python 包来获取一些用于文本分类的嵌入。

常规版本目前可以使用：

import openai
from openai.embeddings_utils import get_embedding

def add_embeddings(data:pd.DataFrame) -> pd.DataFrame:
  data["embedding"] = data['target_column'].apply(lambda x: get_embedding(x, engine="text-embedding-ada-002"))

  return data

但我希望转向异步方法，因为从长远来看它会更有效。我尝试了一些不同的事情，但我当前的版本如下所示：

import asyncio
import openai
from openai.embeddings_utils import aget_embedding

async def add_embeddings(data: pd.DataFrame) -> pd.DataFrame:
  data["embedding"] = await asyncio.gather(*(aget_embedding(x, engine="text-embedding-ada-002") for x in data['target_column']))

  return data

运行该函数时（即

data = add_embeddings(data)

），异步版本返回一个协程对象，而不是预期的文本嵌入。我假设我错过了一些简单的东西，但我对异步的经验非常基本。

就上下文而言，我正在使用 Azure Databricks 并尝试利用以下答案：

• 我可以异步执行“apply”到pandas数据帧中的函数吗？
• 将异步函数应用于 pandas 数据帧的最快方法
• await 在 for 循环中如何工作？

Answer 1

当我尝试运行异步aget_text_embedding时，我遇到了同样的问题，但是使用OllamaEmbedding。当我添加

import nest_asyncio

nest_asyncio.apply()

它跑了。然而，值得注意的是，尽管执行看起来是“并发”的，但当我将执行时间与常规 get_text_embedding 进行比较时，它甚至比 aget 版本还要长。