Skip to content

Commit 4995c00

Browse files
authored
fix: rename PaLM2TextEmbeddingGenerator.predict output columns to be backward compatible (#561)
1 parent 4833908 commit 4995c00

File tree

3 files changed

+527
-560
lines changed

3 files changed

+527
-560
lines changed

bigframes/ml/llm.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
_GEMINI_PRO_ENDPOINT = "gemini-pro"
4545

4646
_ML_GENERATE_TEXT_STATUS = "ml_generate_text_status"
47-
_ML_EMBED_TEXT_STATUS = "ml_generate_embedding_status"
47+
_ML_EMBED_TEXT_STATUS = "ml_embed_text_status"
4848

4949

5050
@log_adapter.class_logger
@@ -390,6 +390,13 @@ def predict(self, X: Union[bpd.DataFrame, bpd.Series]) -> bpd.DataFrame:
390390
}
391391

392392
df = self._bqml_model.generate_embedding(X, options)
393+
df = df.rename(
394+
columns={
395+
"ml_generate_embedding_result": "text_embedding",
396+
"ml_generate_embedding_statistics": "statistics",
397+
"ml_generate_embedding_status": _ML_EMBED_TEXT_STATUS,
398+
}
399+
)
393400

394401
if (df[_ML_EMBED_TEXT_STATUS] != "").any():
395402
warnings.warn(

0 commit comments

Comments
 (0)