airflow.providers.google.cloud.hooks.vertex_ai.generative_model¶

This module contains a Google Cloud Vertex AI Generative Model hook.

Classes¶

`GenerativeModelHook`	Hook for Google Cloud Vertex AI Generative Model APIs.
`ExperimentRunHook`	Use the Vertex AI SDK for Python to create and manage your experiment runs.

Module Contents¶

class airflow.providers.google.cloud.hooks.vertex_ai.generative_model.GenerativeModelHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]¶

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseHook

Hook for Google Cloud Vertex AI Generative Model APIs.

get_text_embedding_model(pretrained_model)[source]¶

Return a Model Garden Model object based on Text Embedding.

get_generative_model(pretrained_model, system_instruction=None, generation_config=None, safety_settings=None, tools=None)[source]¶

Return a Generative Model object.

get_eval_task(dataset, metrics, experiment)[source]¶

Return an EvalTask object.

get_cached_context_model(cached_content_name)[source]¶

Return a Generative Model with Cached Context.

text_embedding_model_get_embeddings(prompt, pretrained_model, location, project_id=PROVIDE_PROJECT_ID)[source]¶

Use the Vertex AI PaLM API to generate text embeddings.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
prompt (str) – Required. Inputs or queries that a user or a program gives to the Vertex AI PaLM API, in order to elicit a specific response.
pretrained_model (str) – A pre-trained model optimized for generating text embeddings.

generative_model_generate_content(contents, location, pretrained_model, tools=None, generation_config=None, safety_settings=None, system_instruction=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Use the Vertex AI Gemini Pro foundation model to generate natural language text.

Parameters:

location (str) – Required. The ID of the Google Cloud location that the service belongs to.
project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
contents (list) – Required. The multi-part content of a message that a user or a program gives to the generative model, in order to elicit a specific response.
generation_config (dict | None) – Optional. Generation configuration settings.
safety_settings (dict | None) – Optional. Per request settings for blocking unsafe content.
tools (list | None) – Optional. A list of tools available to the model during evaluation, such as a data store.
system_instruction (str | None) – Optional. An instruction given to the model to guide its behavior.
pretrained_model (str) – Required. Model, supporting prompts with text-only input, including natural language tasks, multi-turn text and code chat, and code generation. It can output text and code.

supervised_fine_tuning_train(source_model, train_dataset, location, tuned_model_display_name=None, validation_dataset=None, epochs=None, adapter_size=None, learning_rate_multiplier=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Use the Supervised Fine Tuning API to create a tuning job.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
source_model (str) – Required. A pre-trained model optimized for performing natural language tasks such as classification, summarization, extraction, content creation, and ideation.
train_dataset (str) – Required. Cloud Storage URI of your training dataset. The dataset must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples.
tuned_model_display_name (str | None) – Optional. Display name of the TunedModel. The name can be up to 128 characters long and can consist of any UTF-8 characters.
validation_dataset (str | None) – Optional. Cloud Storage URI of your training dataset. The dataset must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples.
epochs (int | None) – Optional. To optimize performance on a specific dataset, try using a higher epoch value. Increasing the number of epochs might improve results. However, be cautious about over-fitting, especially when dealing with small datasets. If over-fitting occurs, consider lowering the epoch number.
adapter_size (Literal[1, 4, 8, 16] | None) – Optional. Adapter size for tuning.
learning_rate_multiplier (float | None) – Optional. Multiplier for adjusting the default learning rate.

count_tokens(contents, location, pretrained_model, project_id=PROVIDE_PROJECT_ID)[source]¶

Use the Vertex AI Count Tokens API to calculate the number of input tokens before sending a request to the Gemini API.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
contents (list) – Required. The multi-part content of a message that a user or a program gives to the generative model, in order to elicit a specific response.
pretrained_model (str) – Required. Model, supporting prompts with text-only input, including natural language tasks, multi-turn text and code chat, and code generation. It can output text and code.

run_evaluation(pretrained_model, eval_dataset, metrics, experiment_name, experiment_run_name, prompt_template, location, generation_config=None, safety_settings=None, system_instruction=None, tools=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Use the Rapid Evaluation API to evaluate a model.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
pretrained_model (str) – Required. A pre-trained model optimized for performing natural language tasks such as classification, summarization, extraction, content creation, and ideation.
eval_dataset (dict) – Required. A fixed dataset for evaluating a model against. Adheres to Rapid Evaluation API.
metrics (list) – Required. A list of evaluation metrics to be used in the experiment. Adheres to Rapid Evaluation API.
experiment_name (str) – Required. The name of the evaluation experiment.
experiment_run_name (str) – Required. The specific run name or ID for this experiment.
prompt_template (str) – Required. The template used to format the model’s prompts during evaluation. Adheres to Rapid Evaluation API.
generation_config (dict | None) – Optional. A dictionary containing generation parameters for the model.
safety_settings (dict | None) – Optional. A dictionary specifying harm category thresholds for blocking model outputs.
system_instruction (str | None) – Optional. An instruction given to the model to guide its behavior.
tools (list | None) – Optional. A list of tools available to the model during evaluation, such as a data store.

create_cached_content(model_name, location, ttl_hours=1, system_instruction=None, contents=None, display_name=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Create CachedContent to reduce the cost of requests that contain repeat content with high input token counts.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
model_name (str) – Required. The name of the publisher model to use for cached content.
system_instruction (Any | None) – Developer set system instruction.
contents (list[Any] | None) – The content to cache.
ttl_hours (float) – The TTL for this resource in hours. The expiration time is computed: now + TTL. Defaults to one hour.
display_name (str | None) – The user-generated meaningful display name of the cached content

generate_from_cached_content(location, cached_content_name, contents, generation_config=None, safety_settings=None, project_id=PROVIDE_PROJECT_ID)[source]¶

Generate a response from CachedContent.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
cached_content_name (str) – Required. The name of the cached content resource.
contents (list) – Required. The multi-part content of a message that a user or a program gives to the generative model, in order to elicit a specific response.
generation_config (dict | None) – Optional. Generation configuration settings.
safety_settings (dict | None) – Optional. Per request settings for blocking unsafe content.

class airflow.providers.google.cloud.hooks.vertex_ai.generative_model.ExperimentRunHook(gcp_conn_id='google_cloud_default', impersonation_chain=None, **kwargs)[source]¶

Bases: airflow.providers.google.common.hooks.base_google.GoogleBaseHook

Use the Vertex AI SDK for Python to create and manage your experiment runs.

delete_experiment_run(experiment_run_name, experiment_name, location, project_id=PROVIDE_PROJECT_ID, delete_backing_tensorboard_run=False)[source]¶

Delete experiment run from the experiment.

Parameters:

project_id (str) – Required. The ID of the Google Cloud project that the service belongs to.
location (str) – Required. The ID of the Google Cloud location that the service belongs to.
experiment_name (str) – Required. The name of the evaluation experiment.
experiment_run_name (str) – Required. The specific run name or ID for this experiment.
delete_backing_tensorboard_run (bool) – Whether to delete the backing Vertex AI TensorBoard run that stores time series metrics for this run.