apache_beam.ml.transforms.embeddings.open_ai module

class apache_beam.ml.transforms.embeddings.open_ai.OpenAITextEmbeddings(model_name: str, columns: list[str], api_key: str | None = None, organization: str | None = None, dimensions: int | None = None, user: str | None = None, max_batch_size: int | None = None, **kwargs)[source]

Bases: EmbeddingsManager

Embedding Config for OpenAI Text Embedding models. Text Embeddings are generated for a batch of text using the OpenAI API.

Parameters:

  • model_name – Name of the OpenAI embedding model

  • columns – The columns where the embeddings will be stored in the output

  • api_key – OpenAI API key

  • organization – OpenAI organization ID

  • dimensions – Specific embedding dimensions to use (if model supports it)

  • user – End-user identifier for tracking and rate limit calculations

  • max_batch_size – Maximum batch size for requests to OpenAI API

get_model_handler() → RemoteModelHandler[source]
get_ptransform_for_processing(**kwargs) → PTransform[source]