diff --git a/notebooks/official/prediction/llm_streaming_prediction.ipynb b/notebooks/official/prediction/llm_streaming_prediction.ipynb
new file mode 100644
index 000000000..06fb61b3d
--- /dev/null
+++ b/notebooks/official/prediction/llm_streaming_prediction.ipynb
@@ -0,0 +1,787 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "copyright"
+ },
+ "outputs": [],
+ "source": [
+ "# Copyright 2023 Google LLC\n",
+ "#\n",
+ "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# https://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "title:generic,gcp"
+ },
+ "source": [
+ "# Vertex AI LLM and streaming prediction\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " View on GitHub\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " Run in Colab\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " \n",
+ " Open in Vertex AI Workbench\n",
+ " \n",
+ " | \n",
+ "
\n",
+ "
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "overview:mlops"
+ },
+ "source": [
+ "## Overview\n",
+ "\n",
+ "\n",
+ "This tutorial demonstrates how to use Vertex AI LLM for making streaming predictions on large language models.\n",
+ "\n",
+ "Learn more about [Vertex AI Language Models](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.language_models.TextGenerationModel#vertexai_language_models_TextGenerationModel_predict_streaming)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "objective:mlops,stage2,get_started_automl_training"
+ },
+ "source": [
+ "### Objective\n",
+ "\n",
+ "In this tutorial, you learn how to use Vertex AI LLM to download pretrained LLM model, make predictions and finetuning the model.\n",
+ "\n",
+ "This tutorial uses the following Google Cloud ML services:\n",
+ "\n",
+ "- `Vertex AI LLM`\n",
+ "- `Vertex AI Prediction`\n",
+ "\n",
+ "The steps performed include:\n",
+ "\n",
+ "- Load a pretrained text generation model.\n",
+ "- Make a non-streaming prediction\n",
+ "- Load a pretrained text generation model, which supports streaming.\n",
+ "- Make a streaming prediction\n",
+ "- Load a pretrained chat model.\n",
+ "- Do a local interactive chat session.\n",
+ "- Do a batch prediction with a text generation model.\n",
+ "- Do a batch prediction with a text embedding model."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "dataset:flowers,icn"
+ },
+ "source": [
+ "### Model\n",
+ "\n",
+ "The pre-trained models used for this tutorial are from Vertex AI LLM repository. The models used are:\n",
+ "\n",
+ "- text-bison\n",
+ "- chat-bison\n",
+ "- text-embedding-gecko"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "fb3451ce8e47"
+ },
+ "source": [
+ "### Costs\n",
+ "This tutorial uses billable components of Google Cloud:\n",
+ "\n",
+ "- Vertex AI\n",
+ "- Cloud Storage\n",
+ "\n",
+ "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "install_mlops"
+ },
+ "source": [
+ "## Installations\n",
+ "\n",
+ "Install the packages required for executing this notebook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "NQ6SdFxji8ks"
+ },
+ "outputs": [],
+ "source": [
+ "! pip3 install --upgrade --quiet google-cloud-aiplatform"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "restart"
+ },
+ "source": [
+ "### Colab only: Uncomment the following cell to restart the kernel"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "D-ZBOjErv5mM"
+ },
+ "outputs": [],
+ "source": [
+ "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+ "# import IPython\n",
+ "\n",
+ "# app = IPython.Application.instance()\n",
+ "# app.kernel.do_shutdown(True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "yfEglUHQk9S3"
+ },
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "### Set your project ID\n",
+ "\n",
+ "**If you don't know your project ID**, try the following:\n",
+ "* Run `gcloud config list`.\n",
+ "* Run `gcloud projects list`.\n",
+ "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "set_project_id"
+ },
+ "outputs": [],
+ "source": [
+ "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
+ "\n",
+ "# Set the project id\n",
+ "! gcloud config set project {PROJECT_ID} --quiet"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "region"
+ },
+ "source": [
+ "#### Region\n",
+ "\n",
+ "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "region"
+ },
+ "outputs": [],
+ "source": [
+ "REGION = \"us-central1\" # @param {type: \"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "gcp_authenticate"
+ },
+ "source": [
+ "### Authenticate your Google Cloud account\n",
+ "\n",
+ "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
+ "\n",
+ "**1. Vertex AI Workbench**\n",
+ "* Do nothing as you are already authenticated.\n",
+ "\n",
+ "**2. Local JupyterLab instance, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ce6043da7b33"
+ },
+ "outputs": [],
+ "source": [
+ "# ! gcloud auth login"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0367eac06a10"
+ },
+ "source": [
+ "**3. Colab, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "21ad4dbb4a61"
+ },
+ "outputs": [],
+ "source": [
+ "# from google.colab import auth\n",
+ "# auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c13224697bfb"
+ },
+ "source": [
+ "**4. Service account or other**\n",
+ "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "bucket:mbsdk"
+ },
+ "source": [
+ "### Create a Cloud Storage bucket\n",
+ "\n",
+ "Create a storage bucket to store intermediate artifacts such as datasets."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "bucket"
+ },
+ "outputs": [],
+ "source": [
+ "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "autoset_bucket"
+ },
+ "source": [
+ "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "91c46850b49b"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "setup_vars"
+ },
+ "source": [
+ "### Set up variables\n",
+ "\n",
+ "Next, set up some variables used throughout the tutorial.\n",
+ "### Import libraries and define constants"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "8696271b314f"
+ },
+ "outputs": [],
+ "source": [
+ "import vertexai\n",
+ "from vertexai.preview.language_models import TextGenerationModel"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "init_aip:mbsdk"
+ },
+ "source": [
+ "### Initialize Vertex AI SDK for Python\n",
+ "\n",
+ "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "oiraqRwdooE0"
+ },
+ "outputs": [],
+ "source": [
+ "vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "YCriCYpOjZv0"
+ },
+ "source": [
+ "## Text generation\n",
+ "\n",
+ "First, you load from Vertex AI LLM, the text-bison model. Once loaded, you make a prediction."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "fkAkKn-ijJR5"
+ },
+ "outputs": [],
+ "source": [
+ "model = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
+ "\n",
+ "print(\n",
+ " model.predict(\n",
+ " \"What is the best recipe for banana bread? Recipe:\",\n",
+ " # \"Brainstorm some ideas combining VR and fitness:\",\n",
+ " # Optional:\n",
+ " # max_output_tokens=128,\n",
+ " # temperature=0,\n",
+ " # top_p=1,\n",
+ " # top_k=5,\n",
+ " )\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "vltcd5jb5WGe"
+ },
+ "source": [
+ "## Streaming prediction\n",
+ "\n",
+ "Next, you load from Vertex AI LLM, a version of the text-bison model that support streaming prediction. Once loaded, you make a prediction. Note, how each iteration that model returns the next successive output in the prediction response."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Z72f5GTm1weU"
+ },
+ "outputs": [],
+ "source": [
+ "import datetime\n",
+ "\n",
+ "text_generation_model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
+ "\n",
+ "print(\"Start: \", datetime.datetime.now())\n",
+ "for response in text_generation_model.predict_streaming(\n",
+ " prompt=\"Count to 100\", max_output_tokens=1000\n",
+ "):\n",
+ " print(datetime.datetime.now())\n",
+ " print(response)\n",
+ "print(\"End: \", datetime.datetime.now())"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "wl2AZceWjXoy"
+ },
+ "source": [
+ "## Chat\n",
+ "\n",
+ "Next, you load from Vertex AI LLM, a version of the chat-bison model Once loaded, you do an interactive chat session."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Lco7I7X_Ch2C"
+ },
+ "outputs": [],
+ "source": [
+ "from vertexai.language_models import ChatModel, InputOutputTextPair\n",
+ "\n",
+ "chat_model2 = ChatModel.from_pretrained(\"google/chat-bison@001\")\n",
+ "\n",
+ "chat2 = chat_model2.start_chat(\n",
+ " # Optional:\n",
+ " context=\"My name is Ned. You are my personal assistant. My favorite movies are Lord of the Rings and Hobbit.\",\n",
+ " examples=[\n",
+ " InputOutputTextPair(\n",
+ " input_text=\"Who do you work for?\",\n",
+ " output_text=\"I work for Ned.\",\n",
+ " ),\n",
+ " InputOutputTextPair(\n",
+ " input_text=\"What do I like?\",\n",
+ " output_text=\"Ned likes watching movies.\",\n",
+ " ),\n",
+ " ],\n",
+ ")\n",
+ "\n",
+ "print(chat2.send_message(\"Are my favorite movies based on a book series?\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "G64BrDoxC-K3"
+ },
+ "outputs": [],
+ "source": [
+ "print(chat2.send_message(\"When where these books published?\"))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "tmWzgmkUzgzU"
+ },
+ "source": [
+ "## Text embedding\n",
+ "\n",
+ "Next, you load from Vertex AI LLM, a version of the text-embedding-gecko model, and then make an embedding request."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Eara9dykkJTA"
+ },
+ "outputs": [],
+ "source": [
+ "from vertexai.language_models import TextEmbeddingModel\n",
+ "\n",
+ "model = TextEmbeddingModel.from_pretrained(\"google/textembedding-gecko@001\")\n",
+ "embeddings = model.get_embeddings([\"What is life?\"])\n",
+ "for embedding in embeddings:\n",
+ " vector = embedding.values\n",
+ " print(len(vector))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "7gaFptssn03c"
+ },
+ "source": [
+ "## Batch prediction\n",
+ "\n",
+ "### text-bison model\n",
+ "\n",
+ "Now, you do a batch prediction job with a pretrained text-bison model."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "xHl1XlTpqIFY"
+ },
+ "outputs": [],
+ "source": [
+ "dataset = \"gs://cloud-samples-data/vertex-ai/prediction/llm/test_table.jsonl\"\n",
+ "destination_uri_prefix = f\"{BUCKET_URI}/text-bison@001_\"\n",
+ "! gsutil cp -r gs://cloud-samples-data/vertex-ai/prediction/llm/text-bison@001_/ {destination_uri_prefix}\n",
+ "\n",
+ "\n",
+ "from vertexai.language_models import TextGenerationModel\n",
+ "\n",
+ "text_generation_model = TextGenerationModel.from_pretrained(\"text-bison\")\n",
+ "batch_job_1 = text_generation_model.batch_predict(\n",
+ " dataset=dataset,\n",
+ " destination_uri_prefix=destination_uri_prefix,\n",
+ " model_parameters={},\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "3d8c7fc4a4a7"
+ },
+ "source": [
+ "### text-embedding-gecko model\n",
+ "\n",
+ "Next, you do a batch prediction job with a pretrained text-embedding-gecko model."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "DfdSZT4SoPxf"
+ },
+ "outputs": [],
+ "source": [
+ "dataset = \"gs://cloud-samples-data/vertex-ai/prediction/llm/embedding_input.jsonl\"\n",
+ "\n",
+ "destination_uri_prefix = f\"{BUCKET_URI}/textembedding-gecko@001_\"\n",
+ "\n",
+ "from vertexai.preview.language_models import TextEmbeddingModel\n",
+ "\n",
+ "text_embedding_model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")\n",
+ "batch_job_2 = text_embedding_model.batch_predict(\n",
+ " dataset=dataset,\n",
+ " destination_uri_prefix=destination_uri_prefix,\n",
+ " # Optional:\n",
+ " model_parameters={},\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Dt6Cyr9B0ZOT"
+ },
+ "source": [
+ "## Tuning\n",
+ "\n",
+ "Now, you fine tune with a pretrained text-bison model, and make a prediction with the fine tuned model. Your input data is in JSONL format and stored in a Cloud Storage location."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "WxmUoYfC0aJI"
+ },
+ "outputs": [],
+ "source": [
+ "model3 = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
+ "\n",
+ "model3.list_tuned_model_names()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Bb4HhsLc0ha4"
+ },
+ "outputs": [],
+ "source": [
+ "tuning_job = model3.tune_model(\n",
+ " training_data=\"gs://cloud-samples-data/vertex-ai/prediction/llm/q_a_train_with_context.jsonl\",\n",
+ " # Optional:\n",
+ " train_steps=1,\n",
+ " tuning_job_location=\"europe-west4\",\n",
+ " tuned_model_location=\"us-central1\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Xf7tWnMU03ZY"
+ },
+ "outputs": [],
+ "source": [
+ "tuned_model = tuning_job.get_tuned_model()\n",
+ "\n",
+ "print(tuned_model.predict(\"Tell me some ideas combining VR and fitness:\"))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "7UqoxmAd28QZ"
+ },
+ "outputs": [],
+ "source": [
+ "model3.list_tuned_model_names()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "uB7RC4Sd2-vK"
+ },
+ "outputs": [],
+ "source": [
+ "tuned_model4 = model3.get_tuned_model(\n",
+ " tuned_model_name=model3.list_tuned_model_names()[0]\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "asNjLVy47_Ox"
+ },
+ "outputs": [],
+ "source": [
+ "print(tuned_model4.predict(\"Brainstorm some ideas combining VR and fitness:\"))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Wai6L3fQowJ_"
+ },
+ "source": [
+ "## Tuning from pandas DataFrame\n",
+ "\n",
+ "Now, you fine tune with a pretrained text-bison model, and make a prediction with the fine tuned model. Your input data is an in-memory pandas DataFrame."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "lUsZKDS3nIcJ"
+ },
+ "outputs": [],
+ "source": [
+ "import pandas\n",
+ "\n",
+ "training_data = pandas.DataFrame(\n",
+ " data=[\n",
+ " {\"input_text\": \"Input 1\", \"output_text\": \"Output 1\"},\n",
+ " {\"input_text\": \"Input 2\", \"output_text\": \"Output 2\"},\n",
+ " {\"input_text\": \"Input 3\", \"output_text\": \"Output 3\"},\n",
+ " {\"input_text\": \"Input 4\", \"output_text\": \"Output 4\"},\n",
+ " {\"input_text\": \"Input 5\", \"output_text\": \"Output 5\"},\n",
+ " {\"input_text\": \"Input 6\", \"output_text\": \"Output 6\"},\n",
+ " {\"input_text\": \"Input 7\", \"output_text\": \"Output 7\"},\n",
+ " {\"input_text\": \"Input 8\", \"output_text\": \"Output 8\"},\n",
+ " {\"input_text\": \"Input 9\", \"output_text\": \"Output 9\"},\n",
+ " {\"input_text\": \"Input 10\", \"output_text\": \"Output 10\"},\n",
+ " ]\n",
+ ")\n",
+ "\n",
+ "training_data"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "RMReD76mnAOC"
+ },
+ "outputs": [],
+ "source": [
+ "model4 = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
+ "\n",
+ "tuning_job = model4.tune_model(\n",
+ " training_data=training_data,\n",
+ " # Optional:\n",
+ " train_steps=10,\n",
+ " tuning_job_location=\"europe-west4\",\n",
+ " tuned_model_location=\"us-central1\",\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "KWCaCpTDnL1h"
+ },
+ "outputs": [],
+ "source": [
+ "tuned_model = tuning_job.get_tuned_model()\n",
+ "\n",
+ "print(tuned_model.predict(\"Tell me some ideas combining VR and fitness:\"))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "b1bb76638fb5"
+ },
+ "source": [
+ "## Cleanup"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "c41c153770e6"
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "batch_job_1.delete()\n",
+ "batch_job_2.delete()\n",
+ "\n",
+ "delete_bucket = False\n",
+ "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+ " ! gsutil rm -rf {BUCKET_URI}"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "name": "llm_streaming_prediction.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}