Skip to content

Commit 81cb659

Browse files
committed
update warning methods
1 parent 5e11ea0 commit 81cb659

File tree

1 file changed

+19
-7
lines changed

1 file changed

+19
-7
lines changed

bigframes/ml/llm.py

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@ class TextEmbeddingGenerator(base.RetriableRemotePredictor):
115115
The model for text embedding. Possible values are "text-embedding-005", "text-embedding-004"
116116
or "text-multilingual-embedding-002". text-embedding models returns model embeddings for text inputs.
117117
text-multilingual-embedding models returns model embeddings for text inputs which support over 100 languages.
118-
Default to "text-embedding-004".
118+
If no setting is provided, "text-embedding-004" will be used by
119+
default and a warning will be issued.
119120
session (bigframes.Session or None):
120121
BQ session to create the model. If None, use the global default session.
121122
connection_name (str or None):
@@ -134,11 +135,10 @@ def __init__(
134135
session: Optional[bigframes.Session] = None,
135136
connection_name: Optional[str] = None,
136137
):
137-
self.model_name = model_name
138-
if model_name == "text-embedding-004":
138+
if model_name == _TEXT_EMBEDDING_004_ENDPOINT:
139139
msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING)
140140
warnings.warn(msg, category=FutureWarning, stacklevel=2)
141-
141+
self.model_name = model_name
142142
self.session = session or global_session.get_global_session()
143143
self.connection_name = connection_name
144144

@@ -262,7 +262,8 @@ class MultimodalEmbeddingGenerator(base.RetriableRemotePredictor):
262262
Args:
263263
model_name (str, Default to "multimodalembedding@001"):
264264
The model for multimodal embedding. Can set to "multimodalembedding@001". Multimodal-embedding models returns model embeddings for text, image and video inputs.
265-
Default to "multimodalembedding@001".
265+
If no setting is provided, "multimodalembedding@001" will be used by
266+
default and a warning will be issued.
266267
session (bigframes.Session or None):
267268
BQ session to create the model. If None, use the global default session.
268269
connection_name (str or None):
@@ -279,6 +280,9 @@ def __init__(
279280
):
280281
if not bigframes.options.experiments.blob:
281282
raise NotImplementedError()
283+
if model_name == _MULTIMODAL_EMBEDDING_001_ENDPOINT:
284+
msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING)
285+
warnings.warn(msg, category=FutureWarning, stacklevel=2)
282286
self.model_name = model_name
283287
self.session = session or global_session.get_global_session()
284288
self.connection_name = connection_name
@@ -414,7 +418,8 @@ class GeminiTextGenerator(base.RetriableRemotePredictor):
414418
"gemini-1.5-pro-001", "gemini-1.5-pro-002", "gemini-1.5-flash-001",
415419
"gemini-1.5-flash-002", "gemini-2.0-flash-exp",
416420
"gemini-2.0-flash-lite-001", and "gemini-2.0-flash-001".
417-
Default to "gemini-2.0-flash-001".
421+
If no setting is provided, "gemini-2.0-flash-001" will be used by
422+
default and a warning will be issued.
418423
419424
.. note::
420425
"gemini-2.0-flash-exp", "gemini-1.5-pro-preview-0514" and "gemini-1.5-flash-preview-0514" is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the
@@ -460,6 +465,9 @@ def __init__(
460465
"(https://cloud.google.com/products#product-launch-stages)."
461466
)
462467
warnings.warn(msg, category=exceptions.PreviewWarning)
468+
if model_name == _GEMINI_2_FLASH_001_ENDPOINT:
469+
msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING)
470+
warnings.warn(msg, category=FutureWarning, stacklevel=2)
463471
self.model_name = model_name
464472
self.session = session or global_session.get_global_session()
465473
self.max_iterations = max_iterations
@@ -809,7 +817,8 @@ class Claude3TextGenerator(base.RetriableRemotePredictor):
809817
"claude-3-5-sonnet" is Anthropic's most powerful AI model and maintains the speed and cost of Claude 3 Sonnet, which is a mid-tier model.
810818
"claude-3-opus" is Anthropic's second-most powerful AI model, with strong performance on highly complex tasks.
811819
https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#available-claude-models
812-
Default to "claude-3-sonnet".
820+
If no setting is provided, "claude-3-sonnet" will be used by default
821+
and a warning will be issued.
813822
session (bigframes.Session or None):
814823
BQ session to create the model. If None, use the global default session.
815824
connection_name (str or None):
@@ -827,6 +836,9 @@ def __init__(
827836
session: Optional[bigframes.Session] = None,
828837
connection_name: Optional[str] = None,
829838
):
839+
if model_name == _CLAUDE_3_SONNET_ENDPOINT:
840+
msg = exceptions.format_message(_REMOVE_DEFAULT_MODEL_WARNING)
841+
warnings.warn(msg, category=FutureWarning, stacklevel=2)
830842
self.model_name = model_name
831843
self.session = session or global_session.get_global_session()
832844
self.connection_name = connection_name

0 commit comments

Comments
 (0)