Vertex AI Embeddings for Text によるテキストエンベディングをやってみた

!pip3 install --upgrade google-cloud-aiplatform \
                        google-cloud-storage \
                        'google-cloud-bigquery[pandas]'

PROJECT_ID = "YOUR-PROJECT-ID"
REGION = "asia-northeast1"
!gcloud config set project {PROJECT_ID}

from google.colab import auth
auth.authenticate_user()

import vertexai
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
import pandas as pd
import json
import os

# Set the model name.
MODEL_NAME = "textembedding-gecko-multilingual@latest"

# Set the task_type, text and optional title as the model inputs.
TASK_TYPE = "CLUSTERING"

# Initiate Vertex AI
vertexai.init(project=PROJECT_ID, location=REGION)

import vertexai
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
import pandas as pd
import json
import os

# Set the model name.
MODEL_NAME = "textembedding-gecko-multilingual@latest"

# Set the task_type, text and optional title as the model inputs.
TASK_TYPE = "CLUSTERING"

# Initiate Vertex AI
vertexai.init(project=PROJECT_ID, location=REGION)

import math
from typing import Any, Generator

import pandas as pd
from google.cloud import bigquery

client = bigquery.Client(project=PROJECT_ID)
QUERY_TEMPLATE = """
        SELECT distinct q.id, q.title, q.body
        FROM (SELECT * FROM `bigquery-public-data.stackoverflow.posts_questions` where Score>0 ORDER BY View_Count desc) AS q
        LIMIT {limit} OFFSET {offset};
        """


def query_bigquery_chunks(
    max_rows: int, rows_per_chunk: int, start_chunk: int = 0
) -> Generator[pd.DataFrame, Any, None]:
    for offset in range(start_chunk, max_rows, rows_per_chunk):
        query = QUERY_TEMPLATE.format(limit=rows_per_chunk, offset=offset)
        query_job = client.query(query)
        rows = query_job.result()
        df = rows.to_dataframe()
        yield df

# 100件のレコードを取得
df = next(query_bigquery_chunks(max_rows=100, rows_per_chunk=100))

from google.colab import drive
drive.mount('/content/drive')

!ls /content/drive/MyDrive/

file_name = "test.csv"
filepath = f'/content/drive/MyDrive/{file_name}'
df = pd.read_csv(filepath)

BUCKET_NAME = "your-bucket-name"
file_name = "test.csv"

import io
from google.cloud import storage

# gcsのクライアントを作成
client = storage.Client()

# バケットを取得
bucket = client.bucket(BUCKET_NAME)

# ファイルのblobインスタンスを取得
blob = bucket.blob(file_name)
# byte文字列としてダウンロード
byte_string = blob.download_as_bytes()

# byte文字列をデコード
file = io.BytesIO(byte_string)
df = pd.read_csv(file)

import vertexai
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel

def text_embedding(
  model_name: str, task_type: str, text: str, title: str = "") -> list:
    """Generate text embedding with a Large Language Model."""
    model = TextEmbeddingModel.from_pretrained(model_name)

    text_embedding_input = TextEmbeddingInput(
        task_type=task_type, title=title, text=text)
    embeddings = model.get_embeddings([text_embedding_input])
    return embeddings[0].values

embeddings = []
for v in df.itertuples():
  TEXT = v.body # メインとなるテキスト
  if TASK_TYPE != 'RETRIEVAL_DOCUMENT':
    embedding = text_embedding(model_name=MODEL_NAME, task_type=TASK_TYPE, text=TEXT)
  else: # RETRIEVAL_DOCUMENTの場合はタイトルの入力が可能
    TITLE = v.title # タイトルとなるテキスト
    embedding = text_embedding(model_name=MODEL_NAME, task_type=TASK_TYPE, text=TEXT, title=TITLE)
  embeddings.append(embedding)

pd.DataFrame(embeddings)