Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Nov 13, 2024
1 parent 1d1b51e commit 792e419
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions comps/embeddings/predictionguard/embedding_predictionguard.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import os
import time
from typing import List, Optional, Union

from predictionguard import PredictionGuard

from comps import (
Expand Down Expand Up @@ -72,12 +73,14 @@ async def embedding(
logger.info(res)
return res


async def get_embeddings(text: Union[str, List[str]]) -> List[List[float]]:
texts = [text] if isinstance(text, str) else text
response = client.embeddings.create(model=pg_embedding_model_name, input=texts)["data"]
embed_vector = [response[i]["embedding"] for i in range(len(response))]
return embed_vector


if __name__ == "__main__":
pg_embedding_model_name = os.getenv("PG_EMBEDDING_MODEL_NAME", "bridgetower-large-itm-mlm-itc")
print("Prediction Guard Embedding initialized.")
Expand Down

0 comments on commit 792e419

Please sign in to comment.