fix!: exclude remote models for .register() (#465)

* fix: exclude remote models for .register() * fix mypy

Author: GarrettWu

bigframes/ml/base.py

@@ -90,6 +90,7 @@ def __repr__(self):
  return prettyprinter.pformat(self)
 
 
+# TODO(garrettwu): refactor to reflect the actual property. Now the class contains .register() method.
 class Predictor(BaseEstimator):
  """A BigQuery DataFrames ML Model base class that can be used to predict outputs."""
 

bigframes/ml/llm.py

@@ -48,7 +48,7 @@
 
 
 @log_adapter.class_logger
-class PaLM2TextGenerator(base.Predictor):
+class PaLM2TextGenerator(base.BaseEstimator):
  """PaLM2 text generator LLM model.
 
  Args:
@@ -258,7 +258,7 @@ def to_gbq(self, model_name: str, replace: bool = False) -> PaLM2TextGenerator:
 
 
 @log_adapter.class_logger
-class PaLM2TextEmbeddingGenerator(base.Predictor):
+class PaLM2TextEmbeddingGenerator(base.BaseEstimator):
  """PaLM2 text embedding generator LLM model.
 
  Args:
@@ -418,7 +418,7 @@ def to_gbq(
 
 
 @log_adapter.class_logger
-class GeminiTextGenerator(base.Predictor):
+class GeminiTextGenerator(base.BaseEstimator):
  """Gemini text generator LLM model.
 
  Args:

tests/system/small/ml/test_register.py

@@ -14,6 +14,8 @@
 
 from typing import cast
 
+import pytest
+
 from bigframes.ml import core, imported, linear_model, llm
 
 
@@ -54,19 +56,8 @@ def test_linear_reg_register_with_params(
 def test_palm2_text_generator_register(
  ephemera_palm2_text_generator_model: llm.PaLM2TextGenerator,
 ):
- model = ephemera_palm2_text_generator_model
- model.register()
-
- model_name = "bigframes_" + cast(
- str, cast(core.BqmlModel, model._bqml_model).model.model_id
- )
- # Only registered model contains the field, and the field includes project/dataset. Here only check model_id.
- assert (
- model_name[:63] # truncated
- in cast(core.BqmlModel, model._bqml_model).model.training_runs[-1][
- "vertexAiModelId"
- ]
- )
+ with pytest.raises(AttributeError):
+ ephemera_palm2_text_generator_model.register() # type: ignore
 
 
 def test_imported_tensorflow_register(