File tree Expand file tree Collapse file tree 3 files changed +4
-4
lines changed Expand file tree Collapse file tree 3 files changed +4
-4
lines changed Original file line number Diff line number Diff line change 17
17
# [START generativeaionvertexai_embedding_code_retrieval]
18
18
from vertexai .language_models import TextEmbeddingInput , TextEmbeddingModel
19
19
20
- MODEL_NAME = "text-embedding-preview-0815 "
20
+ MODEL_NAME = "text-embedding-005 "
21
21
DIMENSIONALITY = 256
22
22
23
23
24
24
def embed_text (
25
25
texts : list [str ] = ["Retrieve a function that adds two numbers" ],
26
26
task : str = "CODE_RETRIEVAL_QUERY" ,
27
- model_name : str = "text-embedding-preview-0815 " ,
27
+ model_name : str = "text-embedding-005 " ,
28
28
dimensionality : int | None = 256 ,
29
29
) -> list [list [float ]]:
30
30
"""Embeds texts with a pre-trained, foundational model."""
Original file line number Diff line number Diff line change @@ -32,7 +32,7 @@ def embed_text() -> list[list[float]]:
32
32
# The task type for embedding. Check the available tasks in the model's documentation.
33
33
task = "RETRIEVAL_DOCUMENT"
34
34
35
- model = TextEmbeddingModel .from_pretrained ("text-embedding-004 " )
35
+ model = TextEmbeddingModel .from_pretrained ("text-embedding-005 " )
36
36
inputs = [TextEmbeddingInput (text , task ) for text in texts ]
37
37
kwargs = dict (output_dimensionality = dimensionality ) if dimensionality else {}
38
38
embeddings = model .get_embeddings (inputs , ** kwargs )
Original file line number Diff line number Diff line change 21
21
22
22
def tune_embedding_model (
23
23
api_endpoint : str ,
24
- base_model_name : str = "text-embedding-004 " ,
24
+ base_model_name : str = "text-embedding-005 " ,
25
25
corpus_path : str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl" ,
26
26
queries_path : str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl" ,
27
27
train_label_path : str = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv" ,
You can’t perform that action at this time.
0 commit comments