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(