move file

notebooks/community/generative_ai/text_embedding_model.ipynb

@@ -0,0 +1,212 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3wy92L3hM7er"
+      },
+      "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": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HD0LnJ0mGjEn"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install git+https://github.com/googleapis/python-aiplatform.git@copybara_559287614\n",
+        "# Remember to restart your runtime!"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "QyhbYch5o-jh"
+      },
+      "outputs": [],
+      "source": [
+        "from google.colab import auth as google_auth\n",
+        "\n",
+        "google_auth.authenticate_user()\n",
+        "import vertexai\n",
+        "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n",
+        "\n",
+        "# Note: when publishing, will make the PROJECT_ID empty\n",
+        "PROJECT_ID = \"cloud-llm-preview1\"  # @param {type:\"string\"}\n",
+        "TASK_TYPE = \"RETRIEVAL_QUERY\"  # @param [\"RETRIEVAL_QUERY\", \"RETRIEVAL_DOCUMENT\", \"SEMANTIC_SIMILARITY\", \"CLASSIFICATION\", \"CLUSTERING\"]\n",
+        "TITLE = \"\"  # @param {type:\"string\"}\n",
+        "TEXT = \"Embedding text.\"  # @param {type:\"string\"}\n",
+        "MODEL_NAME = \"textembedding-gecko@latest\"  # @param [\"textembedding-gecko@latest\", \"textembedding-gecko-multilingual@latest\"]\n",
+        "\n",
+        "if not PROJECT_ID:\n",
+        "    raise ValueError(\"Please set your PROJECT_ID.\")\n",
+        "if not TASK_TYPE:\n",
+        "    raise ValueError(\"Please set TASK_TYPE.\")\n",
+        "if not TEXT:\n",
+        "    raise ValueError(\"Please set TEXT.\")\n",
+        "if not MODEL_NAME:\n",
+        "    raise ValueError(\"Please set MODEL_NAME.\")\n",
+        "if TITLE and TASK_TYPE != \"RETRIEVAL_DOCUMENT\":\n",
+        "    raise ValueError(\n",
+        "        \"Title can only be provided if the task_type is RETRIEVAL_DOCUMENT\"\n",
+        "    )\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=\"us-central1\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "AW1sx9OBCeTL"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "768\n"
+          ]
+        }
+      ],
+      "source": [
+        "def text_embedding(\n",
+        "    model_name: str,\n",
+        "    task_type: str,\n",
+        "    text: str,\n",
+        "    title: str = \"\",\n",
+        ") -> list:\n",
+        "    \"\"\"Text embedding with a Large Language Model.\"\"\"\n",
+        "    model = TextEmbeddingModel.from_pretrained(model_name)\n",
+        "\n",
+        "    text_embedding_input = TextEmbeddingInput(\n",
+        "        task_type=task_type, title=title, text=text\n",
+        "    )\n",
+        "    embeddings = model.get_embeddings([text_embedding_input])\n",
+        "    return embeddings[0].values\n",
+        "\n",
+        "\n",
+        "embedding = text_embedding(\n",
+        "    model_name=MODEL_NAME, task_type=TASK_TYPE, text=TEXT, title=TITLE\n",
+        ")\n",
+        "print(len(embedding))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "I-07lw47qjq7"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "-0.8425277400092455\n",
+            "-0.5084050247105552\n",
+            "-0.6567934176709969\n",
+            "0.3477527925315371\n",
+            "0.06128880959295202\n",
+            "0.283250759935072\n"
+          ]
+        }
+      ],
+      "source": [
+        "# End of public demo\n",
+        "# The following is to test the correctness of the above code.\n",
+        "vertexai.init(project=\"cloud-llm-preview1\", location=\"us-central1\")\n",
+        "print(\n",
+        "    sum(\n",
+        "        text_embedding(\n",
+        "            \"textembedding-gecko@latest\",\n",
+        "            \"SEMANTIC_SIMILARITY\",\n",
+        "            title=\"\",\n",
+        "            text=\"simple embedding\",\n",
+        "        )\n",
+        "    )\n",
+        ")\n",
+        "print(\n",
+        "    sum(\n",
+        "        text_embedding(\n",
+        "            \"textembedding-gecko@latest\",\n",
+        "            \"RETRIEVAL_DOCUMENT\",\n",
+        "            title=\"\",\n",
+        "            text=\"simple embedding\",\n",
+        "        )\n",
+        "    )\n",
+        ")\n",
+        "print(\n",
+        "    sum(\n",
+        "        text_embedding(\n",
+        "            \"textembedding-gecko@latest\",\n",
+        "            \"RETRIEVAL_DOCUMENT\",\n",
+        "            title=\"Google\",\n",
+        "            text=\"simple embedding\",\n",
+        "        )\n",
+        "    )\n",
+        ")\n",
+        "print(\n",
+        "    sum(\n",
+        "        text_embedding(\n",
+        "            \"textembedding-gecko-multilingual@latest\",\n",
+        "            \"SEMANTIC_SIMILARITY\",\n",
+        "            title=\"\",\n",
+        "            text=\"simple embedding\",\n",
+        "        )\n",
+        "    )\n",
+        ")\n",
+        "print(\n",
+        "    sum(\n",
+        "        text_embedding(\n",
+        "            \"textembedding-gecko-multilingual@latest\",\n",
+        "            \"RETRIEVAL_DOCUMENT\",\n",
+        "            title=\"\",\n",
+        "            text=\"simple embedding\",\n",
+        "        )\n",
+        "    )\n",
+        ")\n",
+        "print(\n",
+        "    sum(\n",
+        "        text_embedding(\n",
+        "            \"textembedding-gecko-multilingual@latest\",\n",
+        "            \"RETRIEVAL_DOCUMENT\",\n",
+        "            title=\"Google\",\n",
+        "            text=\"simple embedding\",\n",
+        "        )\n",
+        "    )\n",
+        ")"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "text_embedding_model.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}