Skip to content

Pass in keyword arguments in embedding utility functions #405

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hallacy merged 2 commits into from
Apr 19, 2023
Merged

Pass in keyword arguments in embedding utility functions #405

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8e2e0b5
Pass in kwargs in embedding util functions
tuliren Apr 19, 2023
0e57f61
Remove debug code
tuliren Apr 19, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions openai/embeddings_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,51 +15,51 @@


@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
def get_embedding(text: str, engine="text-similarity-davinci-001") -> List[float]:
def get_embedding(text: str, engine="text-similarity-davinci-001", **kwargs) -> List[float]:

# replace newlines, which can negatively affect performance.
text = text.replace("\n", " ")

return openai.Embedding.create(input=[text], engine=engine)["data"][0]["embedding"]
return openai.Embedding.create(input=[text], engine=engine, **kwargs)["data"][0]["embedding"]


@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
async def aget_embedding(
text: str, engine="text-similarity-davinci-001"
text: str, engine="text-similarity-davinci-001", **kwargs
) -> List[float]:

# replace newlines, which can negatively affect performance.
text = text.replace("\n", " ")

return (await openai.Embedding.acreate(input=[text], engine=engine))["data"][0][
return (await openai.Embedding.acreate(input=[text], engine=engine, **kwargs))["data"][0][
"embedding"
]


@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
def get_embeddings(
list_of_text: List[str], engine="text-similarity-babbage-001"
list_of_text: List[str], engine="text-similarity-babbage-001", **kwargs
) -> List[List[float]]:
assert len(list_of_text) <= 2048, "The batch size should not be larger than 2048."

# replace newlines, which can negatively affect performance.
list_of_text = [text.replace("\n", " ") for text in list_of_text]

data = openai.Embedding.create(input=list_of_text, engine=engine).data
data = openai.Embedding.create(input=list_of_text, engine=engine, **kwargs).data
data = sorted(data, key=lambda x: x["index"]) # maintain the same order as input.
return [d["embedding"] for d in data]


@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6))
async def aget_embeddings(
list_of_text: List[str], engine="text-similarity-babbage-001"
list_of_text: List[str], engine="text-similarity-babbage-001", **kwargs
) -> List[List[float]]:
assert len(list_of_text) <= 2048, "The batch size should not be larger than 2048."

# replace newlines, which can negatively affect performance.
list_of_text = [text.replace("\n", " ") for text in list_of_text]

data = (await openai.Embedding.acreate(input=list_of_text, engine=engine)).data
data = (await openai.Embedding.acreate(input=list_of_text, engine=engine, **kwargs)).data
data = sorted(data, key=lambda x: x["index"]) # maintain the same order as input.
return [d["embedding"] for d in data]

Expand Down