Membuat Embedding untuk Pengambilan Kode

Contoh ini menunjukkan cara menggunakan model embedding teks Vertex AI untuk menghitung embedding bagi blok kode dan kueri untuk tugas pengambilan kode.

Mempelajari lebih lanjut

Untuk dokumentasi mendetail yang menyertakan contoh kode ini, lihat artikel berikut:

Contoh kode

Python

Sebelum mencoba contoh ini, ikuti petunjuk penyiapan Python di Panduan memulai Vertex AI menggunakan library klien. Untuk mengetahui informasi selengkapnya, lihat Dokumentasi referensi API Python Vertex AI.

Untuk melakukan autentikasi ke Vertex AI, siapkan Kredensial Default Aplikasi. Untuk mengetahui informasi selengkapnya, baca Menyiapkan autentikasi untuk lingkungan pengembangan lokal. 

from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel

MODEL_NAME = "text-embedding-005"
DIMENSIONALITY = 256


def embed_text(
    texts: list[str] = ["Retrieve a function that adds two numbers"],
    task: str = "CODE_RETRIEVAL_QUERY",
    model_name: str = "text-embedding-005",
    dimensionality: int | None = 256,
) -> list[list[float]]:
    """Embeds texts with a pre-trained, foundational model."""
    model = TextEmbeddingModel.from_pretrained(model_name)
    inputs = [TextEmbeddingInput(text, task) for text in texts]
    kwargs = dict(output_dimensionality=dimensionality) if dimensionality else {}
    embeddings = model.get_embeddings(inputs, **kwargs)
    # Example response:
    # [[0.025890009477734566, -0.05553026497364044, 0.006374752148985863,...],
    return [embedding.values for embedding in embeddings]


if __name__ == "__main__":
    # Embeds code block with a pre-trained, foundational model.
    # Using this function to calculate the embedding for corpus.
    texts = ["Retrieve a function that adds two numbers"]
    task = "CODE_RETRIEVAL_QUERY"
    code_block_embeddings = embed_text(
        texts=texts, task=task, model_name=MODEL_NAME, dimensionality=DIMENSIONALITY
    )

    # Embeds code retrieval with a pre-trained, foundational model.
    # Using this function to calculate the embedding for query.
    texts = [
        "def func(a, b): return a + b",
        "def func(a, b): return a - b",
        "def func(a, b): return (a ** 2 + b ** 2) ** 0.5",
    ]
    task = "RETRIEVAL_DOCUMENT"
    code_query_embeddings = embed_text(
        texts=texts, task=task, model_name=MODEL_NAME, dimensionality=DIMENSIONALITY
    )

Langkah berikutnya

Untuk menelusuri dan memfilter contoh kode untuk produk Google Cloud lainnya, lihat Google Cloud browser contoh.