From c81298f73dd4f1296cecde461991507974f1aae3 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Wed, 1 May 2024 19:57:42 +0000 Subject: [PATCH] fix: use explicit session in `PaLM2TextGenerator` This is to avoid running into conflict with the global session which may have different options (e.g. location, connection etc.) set. --- bigframes/ml/llm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 9fa9a9acd0..4b07524194 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -144,7 +144,7 @@ def _from_bq( kwargs: dict = {} last_fitting = model.training_runs[-1]["trainingOptions"] - dummy_text_generator = cls() + dummy_text_generator = cls(session=session) for bf_param, _ in dummy_text_generator.__dict__.items(): bqml_param = _BQML_PARAMS_MAPPING.get(bf_param) if bqml_param in last_fitting: