Vertex AI SDK

The vertexai module.

vertexai.init(*, project: Optional[str] = None, location: Optional[str] = None, experiment: Optional[str] = None, experiment_description: Optional[str] = None, experiment_tensorboard: Optional[Union[google.cloud.aiplatform.tensorboard.tensorboard_resource.Tensorboard, str]] = None, staging_bucket: Optional[str] = None, credentials: Optional[google.auth.credentials.Credentials] = None, encryption_spec_key_name: Optional[str] = None, network: Optional[str] = None)

Updates common initialization parameters with provided options.

  • Parameters

    • project (str) – The default project to use when making API calls.

    • location (str) – The default location to use when making API calls. If not set defaults to us-central-1.

    • experiment (str) – Optional. The experiment name.

    • experiment_description (str) – Optional. The description of the experiment.

    • experiment_tensorboard (Union[str, *[tensorboard_resource.Tensorboard](../aiplatform/services.md#google.cloud.aiplatform.Tensorboard)]*) – Optional. The Vertex AI TensorBoard instance, Tensorboard resource name, or Tensorboard resource ID to use as a backing Tensorboard for the provided experiment.

      Example tensorboard resource name format: “projects/123/locations/us-central1/tensorboards/456”

      If experiment_tensorboard is provided and experiment is not, the provided experiment_tensorboard will be set as the global Tensorboard. Any subsequent calls to aiplatform.init() with experiment and without experiment_tensorboard will automatically assign the global Tensorboard to the experiment.

    • staging_bucket (str) – The default staging bucket to use to stage artifacts when making API calls. In the form gs://…

    • credentials (google.auth.credentials.Credentials) – The default custom credentials to use when making API calls. If not provided credentials will be ascertained from the environment.

    • encryption_spec_key_name (Optional[str]) – Optional. The Cloud KMS resource identifier of the customer managed encryption key used to protect a resource. Has the form: projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key. The key needs to be in the same region as where the compute resource is created.

      If set, this resource and all sub-resources will be secured by this key.

    • network (str) – Optional. The full name of the Compute Engine network to which jobs and resources should be peered. E.g. “projects/12345/global/networks/myVPC”. Private services access must already be configured for the network. If specified, all eligible jobs and resources created will be peered with this VPC.

  • Raises

    ValueError – If experiment_description is provided but experiment is not.

Classes for working with language models.

class vertexai.language_models.CodeChatModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._ChatModelBase

CodeChatModel represents a model that is capable of completing code.

Examples

code_chat_model = CodeChatModel.from_pretrained(”codechat-bison@001”)

code_chat = code_chat_model.start_chat(

max_output_tokens=128,
temperature=0.2,

)

code_chat.send_message(“Please help write a function to calculate the min of two numbers”)

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

start_chat(*, max_output_tokens: int = 128, temperature: float = 0.5)

Starts a chat session with the code chat model.

  • Parameters

    • max_output_tokens – Max length of the output text in tokens.

    • temperature – Controls the randomness of predictions. Range: [0, 1].

  • Returns

    A ChatSession object.

class vertexai.language_models.CodeChatSession(model: vertexai.language_models._language_models.CodeChatModel, max_output_tokens: int = 128, temperature: float = 0.5)

Bases: vertexai.language_models._language_models._ChatSessionBase

CodeChatSession represents a chat session with code chat language model.

Within a code chat session, the model keeps context and remembers the previous converstion.

property message_history(: List[vertexai.language_models._language_models.ChatMessage )

List of previous messages.

send_message(message: str, *, max_output_tokens: Optional[int] = None, temperature: Optional[float] = None)

Sends message to the code chat model and gets a response.

  • Parameters

    • message – Message to send to the model

    • max_output_tokens – Max length of the output text in tokens. Uses the value specified when calling CodeChatModel.start_chat by default.

    • temperature – Controls the randomness of predictions. Range: [0, 1]. Uses the value specified when calling CodeChatModel.start_chat by default.

  • Returns

    A TextGenerationResponse object that contains the text produced by the model.

class vertexai.language_models.CodeGenerationModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._LanguageModel

A language model that generates code.

Examples

Getting answers:

generation_model = CodeGenerationModel.from_pretrained(”code-bison@001”) print(generation_model.predict(

prefix=”Write a function that checks if a year is a leap year.”,

))

completion_model = CodeGenerationModel.from_pretrained(”code-gecko@001”) print(completion_model.predict(

prefix=”def reverse_string(s):”,

))

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

predict(prefix: str, suffix: Optional[str] = '', *, max_output_tokens: int = 128, temperature: float = 0.0)

Gets model response for a single prompt.

  • Parameters

    • prefix – Code before the current point.

    • suffix – Code after the current point.

    • max_output_tokens – Max length of the output text in tokens.

    • temperature – Controls the randomness of predictions. Range: [0, 1].

  • Returns

    A TextGenerationResponse object that contains the text produced by the model.

class vertexai.language_models.InputOutputTextPair(input_text: str, output_text: str)

Bases: object

InputOutputTextPair represents a pair of input and output texts.

class vertexai.language_models.TextEmbedding(values: List[float], _prediction_response: Optional[Any] = None)

Bases: object

Contains text embedding vector.

class vertexai.language_models.TextEmbeddingModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._LanguageModel

TextEmbeddingModel converts text into a vector of floating-point numbers.

Examples:

# Getting embedding:
model = TextEmbeddingModel.from_pretrained("textembedding-gecko@001")
embeddings = model.get_embeddings(["What is life?"])
for embedding in embeddings:
    vector = embedding.values
    print(len(vector))

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

class vertexai.language_models.TextGenerationModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._LanguageModel

TextGenerationModel represents a general language model.

Examples:

# Getting answers:
model = TextGenerationModel.from_pretrained("text-bison@001")
model.predict("What is life?")

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

predict(prompt: str, *, max_output_tokens: int = 128, temperature: float = 0.0, top_k: int = 40, top_p: float = 0.95)

Gets model response for a single prompt.

  • Parameters

    • prompt – Question to ask the model.

    • max_output_tokens – Max length of the output text in tokens.

    • temperature – Controls the randomness of predictions. Range: [0, 1].

    • top_k – The number of highest probability vocabulary tokens to keep for top-k-filtering.

    • top_p – The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Range: [0, 1].

  • Returns

    A TextGenerationResponse object that contains the text produced by the model.

class vertexai.language_models.TextGenerationResponse(text: str, _prediction_response: typing.Any, is_blocked: bool = False, safety_attributes: typing.Dict[str, float] =

Bases: object

TextGenerationResponse represents a response of a language model. .. attribute:: text

The generated text

  • type

    str

is_blocked()

Whether the the request was blocked.

safety_attributes()

Scores for safety attributes. Learn more about the safety attributes here: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/responsible-ai#safety_attribute_descriptions

Classes for working with language models.

class vertexai.preview.language_models.ChatMessage(content: str, author: str)

Bases: object

A chat message.

content()

Content of the message.

author()

Author of the message.

class vertexai.preview.language_models.ChatModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._ChatModelBase

ChatModel represents a language model that is capable of chat.

Examples:

chat_model = ChatModel.from_pretrained("chat-bison@001")

chat = chat_model.start_chat(
    context="My name is Ned. You are my personal assistant. My favorite movies are Lord of the Rings and Hobbit.",
    examples=[
        InputOutputTextPair(
            input_text="Who do you work for?",
            output_text="I work for Ned.",
        ),
        InputOutputTextPair(
            input_text="What do I like?",
            output_text="Ned likes watching movies.",
        ),
    ],
    temperature=0.3,
)

chat.send_message("Do you know any cool events this weekend?")

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

start_chat(*, context: Optional[str] = None, examples: Optional[List[vertexai.language_models._language_models.InputOutputTextPair]] = None, max_output_tokens: int = 128, temperature: float = 0.0, top_k: int = 40, top_p: float = 0.95, message_history: Optional[List[vertexai.language_models._language_models.ChatMessage]] = None)

Starts a chat session with the model.

  • Parameters

    • context – Context shapes how the model responds throughout the conversation. For example, you can use context to specify words the model can or cannot use, topics to focus on or avoid, or the response format or style

    • examples – List of structured messages to the model to learn how to respond to the conversation. A list of InputOutputTextPair objects.

    • max_output_tokens – Max length of the output text in tokens.

    • temperature – Controls the randomness of predictions. Range: [0, 1].

    • top_k – The number of highest probability vocabulary tokens to keep for top-k-filtering. Range: [1, 40]

    • top_p – The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Range: [0, 1].

    • message_history – A list of previously sent and received messages.

  • Returns

    A ChatSession object.

class vertexai.preview.language_models.ChatSession(model: vertexai.language_models._language_models.ChatModel, context: Optional[str] = None, examples: Optional[List[vertexai.language_models._language_models.InputOutputTextPair]] = None, max_output_tokens: int = 128, temperature: float = 0.0, top_k: int = 40, top_p: float = 0.95, message_history: Optional[List[vertexai.language_models._language_models.ChatMessage]] = None)

Bases: vertexai.language_models._language_models._ChatSessionBase

ChatSession represents a chat session with a language model.

Within a chat session, the model keeps context and remembers the previous conversation.

property message_history(: List[vertexai.language_models._language_models.ChatMessage )

List of previous messages.

send_message(message: str, *, max_output_tokens: Optional[int] = None, temperature: Optional[float] = None, top_k: Optional[int] = None, top_p: Optional[float] = None)

Sends message to the language model and gets a response.

  • Parameters

    • message – Message to send to the model

    • max_output_tokens – Max length of the output text in tokens. Uses the value specified when calling ChatModel.start_chat by default.

    • temperature – Controls the randomness of predictions. Range: [0, 1]. Uses the value specified when calling ChatModel.start_chat by default.

    • top_k – The number of highest probability vocabulary tokens to keep for top-k-filtering. Uses the value specified when calling ChatModel.start_chat by default.

    • top_p – The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Range: [0, 1]. Uses the value specified when calling ChatModel.start_chat by default.

  • Returns

    A TextGenerationResponse object that contains the text produced by the model.

class vertexai.preview.language_models.CodeChatModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._ChatModelBase

CodeChatModel represents a model that is capable of completing code.

Examples

code_chat_model = CodeChatModel.from_pretrained(”codechat-bison@001”)

code_chat = code_chat_model.start_chat(

max_output_tokens=128,
temperature=0.2,

)

code_chat.send_message(“Please help write a function to calculate the min of two numbers”)

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

start_chat(*, max_output_tokens: int = 128, temperature: float = 0.5)

Starts a chat session with the code chat model.

  • Parameters

    • max_output_tokens – Max length of the output text in tokens.

    • temperature – Controls the randomness of predictions. Range: [0, 1].

  • Returns

    A ChatSession object.

class vertexai.preview.language_models.CodeChatSession(model: vertexai.language_models._language_models.CodeChatModel, max_output_tokens: int = 128, temperature: float = 0.5)

Bases: vertexai.language_models._language_models._ChatSessionBase

CodeChatSession represents a chat session with code chat language model.

Within a code chat session, the model keeps context and remembers the previous converstion.

property message_history(: List[vertexai.language_models._language_models.ChatMessage )

List of previous messages.

send_message(message: str, *, max_output_tokens: Optional[int] = None, temperature: Optional[float] = None)

Sends message to the code chat model and gets a response.

  • Parameters

    • message – Message to send to the model

    • max_output_tokens – Max length of the output text in tokens. Uses the value specified when calling CodeChatModel.start_chat by default.

    • temperature – Controls the randomness of predictions. Range: [0, 1]. Uses the value specified when calling CodeChatModel.start_chat by default.

  • Returns

    A TextGenerationResponse object that contains the text produced by the model.

class vertexai.preview.language_models.CodeGenerationModel(model_id: str, endpoint_name: Optional[str] = None)

Bases: vertexai.language_models._language_models._LanguageModel

A language model that generates code.

Examples

Getting answers:

generation_model = CodeGenerationModel.from_pretrained(”code-bison@001”) print(generation_model.predict(

prefix=”Write a function that checks if a year is a leap year.”,

))

completion_model = CodeGenerationModel.from_pretrained(”code-gecko@001”) print(completion_model.predict(

prefix=”def reverse_string(s):”,

))

Creates a LanguageModel.

This constructor should not be called directly. Use LanguageModel.from_pretrained(model_name=…) instead.

  • Parameters

    • model_id – Identifier of a Vertex LLM. Example: “text-bison@001”

    • endpoint_name – Vertex Endpoint resource name for the model

classmethod from_pretrained(model_name: str)

Loads a _ModelGardenModel.

  • Parameters

    model_name – Name of the model.

  • Returns

    An instance of a class derieved from _ModelGardenModel.

  • Raises

predict(prefix: str, suffix: Optional[str] = '', *, max_output_tokens: int = 128, temperature: float = 0.0)

Gets model response for a single prompt.

  • Parameters

    • prefix – Code before the current point.

    • suffix – Code after the current point.

    • max_output_tokens – Max length of the output text in tokens.

    • temperature – Controls the randomness of predictions. Range: [0, 1].

  • Returns

    A TextGenerationResponse object that contains the text produced by the model.

class vertexai.preview.language_models.InputOutputTextPair(input_text: str, output_text: str)

Bases: object

InputOutputTextPair represents a pair of input and output texts.

class vertexai.preview.language_models.TextEmbedding(values: List[float], _prediction_response: Optional[Any] = None)

Bases: object

Contains text embedding vector.

vertexai.preview.language_models.TextEmbeddingModel()

alias of vertexai.language_models._language_models._PreviewTextEmbeddingModel

vertexai.preview.language_models.TextGenerationModel()

alias of vertexai.language_models._language_models._PreviewTextGenerationModel

class vertexai.preview.language_models.TextGenerationResponse(text: str, _prediction_response: typing.Any, is_blocked: bool = False, safety_attributes: typing.Dict[str, float] =

Bases: object

TextGenerationResponse represents a response of a language model. .. attribute:: text

The generated text

  • type

    str

is_blocked()

Whether the the request was blocked.

safety_attributes()

Scores for safety attributes. Learn more about the safety attributes here: https://cloud.google.com/vertex-ai/docs/generative-ai/learn/responsible-ai#safety_attribute_descriptions