This repository was archived by the owner on Apr 1, 2026. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 68
Expand file tree
/
Copy pathai.py
More file actions
846 lines (687 loc) · 34.2 KB
/
ai.py
File metadata and controls
846 lines (687 loc) · 34.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import re
import typing
from typing import Dict, Iterable, List, Optional, Sequence, Union
import warnings
from bigframes import dtypes, exceptions, options
from bigframes.core import guid
from bigframes.core.logging import log_adapter
@log_adapter.class_logger
class AIAccessor:
def __init__(self, df, base_bqml=None) -> None:
import bigframes # Import in the function body to avoid circular imports.
import bigframes.dataframe
from bigframes.ml import core as ml_core
self._df: bigframes.dataframe.DataFrame = df
self._base_bqml: ml_core.BaseBqml = base_bqml or ml_core.BaseBqml(df._session)
def filter(
self,
instruction: str,
model,
ground_with_google_search: bool = False,
):
"""
Filters the DataFrame with the semantics of the user instruction.
**Examples:**
>>> import bigframes.pandas as bpd
>>> bpd.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001")
>>> df = bpd.DataFrame({"country": ["USA", "Germany"], "city": ["Seattle", "Berlin"]})
>>> df.ai.filter("{city} is the capital of {country}", model)
country city
1 Germany Berlin
<BLANKLINE>
[1 rows x 2 columns]
Args:
instruction (str):
An instruction on how to filter the data. This value must contain
column references by name, which should be wrapped in a pair of braces.
For example, if you have a column "food", you can refer to this column
in the instructions like:
"The {food} is healthy."
model (bigframes.ml.llm.GeminiTextGenerator):
A GeminiTextGenerator provided by Bigframes ML package.
ground_with_google_search (bool, default False):
Enables Grounding with Google Search for the GeminiTextGenerator model.
When set to True, the model incorporates relevant information from Google
Search results into its responses, enhancing their accuracy and factualness.
Note: Using this feature may impact billing costs. Refer to the pricing
page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models
The default is `False`.
Returns:
bigframes.pandas.DataFrame: DataFrame filtered by the instruction.
Raises:
NotImplementedError: when the AI operator experiment is off.
ValueError: when the instruction refers to a non-existing column, or when no
columns are referred to.
"""
if not options.experiments.ai_operators:
raise NotImplementedError()
answer_col = "answer"
output_schema = {answer_col: "bool"}
result = self.map(
instruction,
model,
output_schema,
ground_with_google_search,
)
return result[result[answer_col]].drop(answer_col, axis=1)
def map(
self,
instruction: str,
model,
output_schema: Dict[str, str] | None = None,
ground_with_google_search: bool = False,
):
"""
Maps the DataFrame with the semantics of the user instruction. The name of the keys in the output_schema parameter carry
semantic meaning, and can be used for information extraction.
**Examples:**
>>> import bigframes.pandas as bpd
>>> bpd.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001")
>>> df = bpd.DataFrame({"ingredient_1": ["Burger Bun", "Soy Bean"], "ingredient_2": ["Beef Patty", "Bittern"]})
>>> df.ai.map("What is the food made from {ingredient_1} and {ingredient_2}? One word only.", model=model, output_schema={"food": "string"})
ingredient_1 ingredient_2 food
0 Burger Bun Beef Patty Burger
<BLANKLINE>
1 Soy Bean Bittern Tofu
<BLANKLINE>
<BLANKLINE>
[2 rows x 3 columns]
>>> import bigframes.pandas as bpd
>>> bpd.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001")
>>> df = bpd.DataFrame({"text": ["Elmo lives at 123 Sesame Street."]})
>>> df.ai.map("{text}", model=model, output_schema={"person": "string", "address": "string"})
text person address
0 Elmo lives at 123 Sesame Street. Elmo 123 Sesame Street
<BLANKLINE>
[1 rows x 3 columns]
Args:
instruction (str):
An instruction on how to map the data. This value must contain
column references by name, which should be wrapped in a pair of braces.
For example, if you have a column "food", you can refer to this column
in the instructions like:
"Get the ingredients of {food}."
model (bigframes.ml.llm.GeminiTextGenerator):
A GeminiTextGenerator provided by Bigframes ML package.
output_schema (Dict[str, str] or None, default None):
The schema used to generate structured output as a bigframes DataFrame. The schema is a string key-value pair of <column_name>:<type>.
Supported types are int64, float64, bool, string, array<type> and struct<column type>. If None, generate string result under the column
"ml_generate_text_llm_result".
ground_with_google_search (bool, default False):
Enables Grounding with Google Search for the GeminiTextGenerator model.
When set to True, the model incorporates relevant information from Google
Search results into its responses, enhancing their accuracy and factualness.
Note: Using this feature may impact billing costs. Refer to the pricing
page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models
The default is `False`.
Returns:
bigframes.pandas.DataFrame: DataFrame with attached mapping results.
Raises:
NotImplementedError: when the AI operator experiment is off.
ValueError: when the instruction refers to a non-existing column, or when no
columns are referred to.
"""
if not options.experiments.ai_operators:
raise NotImplementedError()
import bigframes.dataframe
import bigframes.series
self._validate_model(model)
columns = self._parse_columns(instruction)
for column in columns:
if column not in self._df.columns:
raise ValueError(f"Column {column} not found.")
if ground_with_google_search:
msg = exceptions.format_message(
"Enables Grounding with Google Search may impact billing cost. See pricing "
"details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models"
)
warnings.warn(msg, category=UserWarning)
self._confirm_operation(len(self._df))
df: bigframes.dataframe.DataFrame = self._df[columns].copy()
has_blob_column = False
for column in columns:
if df[column].dtype == dtypes.OBJ_REF_DTYPE:
# Don't cast ObjectRef columns to string
has_blob_column = True
continue
if df[column].dtype != dtypes.STRING_DTYPE:
df[column] = df[column].astype(dtypes.STRING_DTYPE)
user_instruction = self._format_instruction(instruction, columns)
output_instruction = (
"Based on the provided contenxt, answer the following instruction:"
)
if output_schema is None:
output_schema = {"ml_generate_text_llm_result": "string"}
if has_blob_column:
results = typing.cast(
bigframes.series.Series,
model.predict(
df,
prompt=self._make_multimodel_prompt(
df, columns, user_instruction, output_instruction
),
temperature=0.0,
ground_with_google_search=ground_with_google_search,
output_schema=output_schema,
),
)
else:
results = typing.cast(
bigframes.series.Series,
model.predict(
self._make_text_prompt(
df, columns, user_instruction, output_instruction
),
temperature=0.0,
ground_with_google_search=ground_with_google_search,
output_schema=output_schema,
),
)
attach_columns = [results[col] for col, _ in output_schema.items()]
from bigframes.core.reshape.api import concat
return concat([self._df, *attach_columns], axis=1)
def classify(
self,
instruction: str,
model,
labels: Sequence[str],
output_column: str = "result",
ground_with_google_search: bool = False,
):
"""
Classifies the rows of dataframes based on user instruction into the provided labels.
**Examples:**
>>> import bigframes.pandas as bpd
>>> bpd.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001")
>>> df = bpd.DataFrame({
... "feedback_text": [
... "The product is amazing, but the shipping was slow.",
... "I had an issue with my recent bill.",
... "The user interface is very intuitive."
... ],
... })
>>> df.ai.classify("{feedback_text}", model=model, labels=["Shipping", "Billing", "UI"])
feedback_text result
0 The product is amazing, but the shipping was s... Shipping
1 I had an issue with my recent bill. Billing
2 The user interface is very intuitive. UI
<BLANKLINE>
[3 rows x 2 columns]
Args:
instruction (str):
An instruction on how to classify the data. This value must contain
column references by name, which should be wrapped in a pair of braces.
For example, if you have a column "feedback", you can refer to this column
with"{food}".
model (bigframes.ml.llm.GeminiTextGenerator):
A GeminiTextGenerator provided by Bigframes ML package.
labels (Sequence[str]):
A collection of labels (categories). It must contain at least two and at most 20 elements.
Labels are case sensitive. Duplicated labels are not allowed.
output_column (str, default "result"):
The name of column for the output.
ground_with_google_search (bool, default False):
Enables Grounding with Google Search for the GeminiTextGenerator model.
When set to True, the model incorporates relevant information from Google
Search results into its responses, enhancing their accuracy and factualness.
Note: Using this feature may impact billing costs. Refer to the pricing
page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models
The default is `False`.
Returns:
bigframes.pandas.DataFrame: DataFrame with classification result.
Raises:
NotImplementedError: when the AI operator experiment is off.
ValueError: when the instruction refers to a non-existing column, when no
columns are referred to, or when the count of labels does not meet the
requirement.
"""
if not options.experiments.ai_operators:
raise NotImplementedError()
if len(labels) < 2 or len(labels) > 20:
raise ValueError(
f"The number of labels should be between 2 and 20 (inclusive), but {len(labels)} labels are provided."
)
if len(set(labels)) != len(labels):
raise ValueError("There are duplicate labels.")
updated_instruction = f"Based on the user instruction {instruction}, you must provide an answer that must exist in the following list of labels: {labels}"
return self.map(
updated_instruction,
model,
output_schema={output_column: "string"},
ground_with_google_search=ground_with_google_search,
)
def join(
self,
other,
instruction: str,
model,
ground_with_google_search: bool = False,
):
"""
Joines two dataframes by applying the instruction over each pair of rows from
the left and right table.
**Examples:**
>>> import bigframes.pandas as bpd
>>> bpd.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.GeminiTextGenerator(model_name="gemini-2.0-flash-001")
>>> cities = bpd.DataFrame({'city': ['Seattle', 'Ottawa', 'Berlin', 'Shanghai', 'New Delhi']})
>>> continents = bpd.DataFrame({'continent': ['North America', 'Africa', 'Asia']})
>>> cities.ai.join(continents, "{city} is in {continent}", model)
city continent
0 Seattle North America
1 Ottawa North America
2 Shanghai Asia
3 New Delhi Asia
<BLANKLINE>
[4 rows x 2 columns]
Args:
other (bigframes.pandas.DataFrame):
The other dataframe.
instruction (str):
An instruction on how left and right rows can be joined. This value must contain
column references by name. which should be wrapped in a pair of braces.
For example: "The {city} belongs to the {country}".
For column names that are shared between two dataframes, you need to add "left."
and "right." prefix for differentiation. This is especially important when you do
self joins. For example: "The {left.employee_name} reports to {right.employee_name}"
For unique column names, this prefix is optional.
model (bigframes.ml.llm.GeminiTextGenerator):
A GeminiTextGenerator provided by Bigframes ML package.
ground_with_google_search (bool, default False):
Enables Grounding with Google Search for the GeminiTextGenerator model.
When set to True, the model incorporates relevant information from Google
Search results into its responses, enhancing their accuracy and factualness.
Note: Using this feature may impact billing costs. Refer to the pricing
page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models
The default is `False`.
Returns:
bigframes.pandas.DataFrame: The joined dataframe.
Raises:
ValueError if the amount of data that will be sent for LLM processing is larger than max_rows.
"""
if not options.experiments.ai_operators:
raise NotImplementedError()
self._validate_model(model)
columns = self._parse_columns(instruction)
if ground_with_google_search:
msg = exceptions.format_message(
"Enables Grounding with Google Search may impact billing cost. See pricing "
"details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models"
)
warnings.warn(msg, category=UserWarning)
work_estimate = len(self._df) * len(other)
self._confirm_operation(work_estimate)
left_columns = []
right_columns = []
for col in columns:
if col in self._df.columns and col in other.columns:
raise ValueError(f"Ambiguous column reference: {col}")
elif col in self._df.columns:
left_columns.append(col)
elif col in other.columns:
right_columns.append(col)
elif col.startswith("left."):
original_col_name = col[len("left.") :]
if (
original_col_name in self._df.columns
and original_col_name in other.columns
):
left_columns.append(col)
elif original_col_name in self._df.columns:
left_columns.append(col)
instruction = instruction.replace(col, original_col_name)
else:
raise ValueError(f"Column {col} not found")
elif col.startswith("right."):
original_col_name = col[len("right.") :]
if (
original_col_name in self._df.columns
and original_col_name in other.columns
):
right_columns.append(col)
elif original_col_name in other.columns:
right_columns.append(col)
instruction = instruction.replace(col, original_col_name)
else:
raise ValueError(f"Column {col} not found")
else:
raise ValueError(f"Column {col} not found")
if not left_columns:
raise ValueError("No left column references.")
if not right_columns:
raise ValueError("No right column references.")
# Update column references to be compatible with internal naming scheme.
# That is, "left.col" -> "col_left" and "right.col" -> "col_right"
instruction = re.sub(r"(?<!{){left\.(\w+)}(?!})", r"{\1_left}", instruction)
instruction = re.sub(r"(?<!{){right\.(\w+)}(?!})", r"{\1_right}", instruction)
joined_df = self._df.merge(other, how="cross", suffixes=("_left", "_right"))
return joined_df.ai.filter(
instruction,
model,
ground_with_google_search=ground_with_google_search,
).reset_index(drop=True)
def search(
self,
search_column: str,
query: str,
top_k: int,
model,
score_column: Optional[str] = None,
):
"""
Performs AI semantic search on the DataFrame.
** Examples: **
>>> import bigframes.pandas as bpd
>>> import bigframes
>>> bigframes.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005")
>>> df = bpd.DataFrame({"creatures": ["salmon", "sea urchin", "frog", "chimpanzee"]})
>>> df.ai.search("creatures", "monkey", top_k=1, model=model, score_column='distance')
creatures distance
3 chimpanzee 0.635844
<BLANKLINE>
[1 rows x 2 columns]
Args:
search_column:
The name of the column to search from.
query (str):
The search query.
top_k (int):
The number of nearest neighbors to return.
model (TextEmbeddingGenerator):
A TextEmbeddingGenerator provided by Bigframes ML package.
score_column (Optional[str], default None):
The name of the the additional column containning the similarity scores. If None,
this column won't be attached to the result.
Returns:
DataFrame: the DataFrame with the search result.
Raises:
ValueError: when the search_column is not found from the the data frame.
TypeError: when the provided model is not TextEmbeddingGenerator.
"""
if not options.experiments.ai_operators:
raise NotImplementedError()
if search_column not in self._df.columns:
raise ValueError(f"Column `{search_column}` not found")
self._confirm_operation(len(self._df))
import bigframes.ml.llm as llm
if not isinstance(model, llm.TextEmbeddingGenerator):
raise TypeError(f"Expect a text embedding model, but got: {type(model)}")
if top_k < 1:
raise ValueError("top_k must be an integer greater than or equal to 1.")
embedded_df = model.predict(self._df[search_column])
embedded_table = embedded_df.reset_index().to_gbq()
import bigframes.pandas as bpd
embedding_result_column = "ml_generate_embedding_result"
query_df = model.predict(bpd.DataFrame({"query_id": [query]})).rename(
columns={"content": "query_id", embedding_result_column: "embedding"}
)
import bigframes.bigquery as bbq
search_result = (
bbq.vector_search(
base_table=embedded_table,
column_to_search=embedding_result_column,
query=query_df,
top_k=top_k,
# TODO(tswast): set allow_large_results based on Series size.
# If we expect small results, it could be faster to set
# allow_large_results to False.
allow_large_results=True,
)
.rename(columns={"content": search_column})
.set_index("index")
)
search_result.index.name = self._df.index.name
if score_column is not None:
search_result = search_result.rename(columns={"distance": score_column})[
[search_column, score_column]
]
else:
search_result = search_result[[search_column]]
import bigframes.dataframe
return typing.cast(bigframes.dataframe.DataFrame, search_result)
def sim_join(
self,
other,
left_on: str,
right_on: str,
model,
top_k: int = 3,
score_column: Optional[str] = None,
max_rows: int = 1000,
):
"""
Joins two dataframes based on the similarity of the specified columns.
This method uses BigQuery's VECTOR_SEARCH function to match rows on the left side with the rows that have
nearest embedding vectors on the right. In the worst case scenario, the complexity is around O(M * N * log K).
Therefore, this is a potentially expensive operation.
** Examples: **
>>> import bigframes.pandas as bpd
>>> bpd.options.experiments.ai_operators = True
>>> bpd.options.compute.ai_ops_confirmation_threshold = 25
>>> import bigframes.ml.llm as llm
>>> model = llm.TextEmbeddingGenerator(model_name="text-embedding-005")
>>> df1 = bpd.DataFrame({'animal': ['monkey', 'spider']})
>>> df2 = bpd.DataFrame({'animal': ['scorpion', 'baboon']})
>>> df1.ai.sim_join(df2, left_on='animal', right_on='animal', model=model, top_k=1)
animal animal_1
0 monkey baboon
1 spider scorpion
<BLANKLINE>
[2 rows x 2 columns]
Args:
other (DataFrame):
The other data frame to join with.
left_on (str):
The name of the column on left side for the join.
right_on (str):
The name of the column on the right side for the join.
top_k (int, default 3):
The number of nearest neighbors to return.
model (TextEmbeddingGenerator):
A TextEmbeddingGenerator provided by Bigframes ML package.
score_column (Optional[str], default None):
The name of the the additional column containning the similarity scores. If None,
this column won't be attached to the result.
max_rows:
The maximum number of rows allowed to be processed per call. If the result is too large, the method
call will end early with an error.
Returns:
DataFrame: the data frame with the join result.
Raises:
ValueError: when the amount of data to be processed exceeds the specified max_rows.
"""
if not options.experiments.ai_operators:
raise NotImplementedError()
if left_on not in self._df.columns:
raise ValueError(f"Left column {left_on} not found")
if right_on not in self._df.columns:
raise ValueError(f"Right column {right_on} not found")
import bigframes.ml.llm as llm
if not isinstance(model, llm.TextEmbeddingGenerator):
raise TypeError(f"Expect a text embedding model, but got: {type(model)}")
joined_table_rows = len(self._df) * len(other)
if joined_table_rows > max_rows:
raise ValueError(
f"Number of rows that need processing is {joined_table_rows}, which exceeds row limit {max_rows}."
)
if top_k < 1:
raise ValueError("top_k must be an integer greater than or equal to 1.")
work_estimate = len(self._df) * len(other)
self._confirm_operation(work_estimate)
base_table_embedding_column = guid.generate_guid()
base_table = self._attach_embedding(
other, right_on, base_table_embedding_column, model
).to_gbq()
query_table = self._attach_embedding(self._df, left_on, "embedding", model)
import bigframes.bigquery as bbq
join_result = bbq.vector_search(
base_table=base_table,
column_to_search=base_table_embedding_column,
query=query_table,
top_k=top_k,
)
join_result = join_result.drop(
["embedding", base_table_embedding_column], axis=1
)
if score_column is not None:
join_result = join_result.rename(columns={"distance": score_column})
else:
del join_result["distance"]
return join_result
def forecast(
self,
timestamp_column: str,
data_column: str,
*,
model: str = "TimesFM 2.0",
id_columns: Optional[Iterable[str]] = None,
horizon: int = 10,
confidence_level: float = 0.95,
):
"""
Forecast time series at future horizon. Using Google Research's open source TimesFM(https://github.com/google-research/timesfm) model.
.. note::
This product or feature is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the
Service Specific Terms(https://cloud.google.com/terms/service-terms#1). Pre-GA products and features are available "as is"
and might have limited support. For more information, see the launch stage descriptions
(https://cloud.google.com/products#product-launch-stages).
Args:
timestamp_column (str):
A str value that specified the name of the time points column.
The time points column provides the time points used to generate the forecast.
The time points column must use one of the following data types: TIMESTAMP, DATE and DATETIME
data_column (str):
A str value that specifies the name of the data column. The data column contains the data to forecast.
The data column must use one of the following data types: INT64, NUMERIC and FLOAT64
model (str, default "TimesFM 2.0"):
A str value that specifies the name of the model. TimesFM 2.0 is the only supported value, and is the default value.
id_columns (Iterable[str] or None, default None):
An iterable of str value that specifies the names of one or more ID columns. Each ID identifies a unique time series to forecast.
Specify one or more values for this argument in order to forecast multiple time series using a single query.
The columns that you specify must use one of the following data types: STRING, INT64, ARRAY<STRING> and ARRAY<INT64>
horizon (int, default 10):
An int value that specifies the number of time points to forecast. The default value is 10. The valid input range is [1, 10,000].
confidence_level (float, default 0.95):
A FLOAT64 value that specifies the percentage of the future values that fall in the prediction interval.
The default value is 0.95. The valid input range is [0, 1).
Returns:
DataFrame:
The forecast dataframe matches that of the BigQuery AI.FORECAST function.
See: https://cloud.google.com/bigquery/docs/reference/standard-sql/bigqueryml-syntax-ai-forecast
Raises:
ValueError: when referring to a non-existing column.
"""
columns = [timestamp_column, data_column]
if id_columns:
columns += id_columns
for column in columns:
if column not in self._df.columns:
raise ValueError(f"Column `{column}` not found")
options: dict[str, Union[int, float, str, Iterable[str]]] = {
"data_col": data_column,
"timestamp_col": timestamp_column,
"model": model,
"horizon": horizon,
"confidence_level": confidence_level,
}
if id_columns:
options["id_cols"] = id_columns
return self._base_bqml.ai_forecast(input_data=self._df, options=options)
@staticmethod
def _attach_embedding(dataframe, source_column: str, embedding_column: str, model):
result_df = dataframe.copy()
embeddings = model.predict(dataframe[source_column])[
"ml_generate_embedding_result"
]
result_df[embedding_column] = embeddings
return result_df
@staticmethod
def _make_multimodel_prompt(
prompt_df, columns, user_instruction: str, output_instruction: str
):
prompt = [f"{output_instruction}\n{user_instruction}\nContext: "]
for col in columns:
prompt.extend([f"{col} is ", prompt_df[col]])
return prompt
@staticmethod
def _make_text_prompt(
prompt_df, columns, user_instruction: str, output_instruction: str
):
prompt_df["prompt"] = f"{output_instruction}\n{user_instruction}\nContext: "
# Combine context from multiple columns.
for col in columns:
prompt_df["prompt"] += f"{col} is `" + prompt_df[col] + "`\n"
return prompt_df["prompt"]
@staticmethod
def _parse_columns(instruction: str) -> List[str]:
"""Extracts column names enclosed in curly braces from the user instruction.
For example, _parse_columns("{city} is in {continent}") == ["city", "continent"]
"""
columns = re.findall(r"(?<!{)\{(?!{)(.*?)\}(?!\})", instruction)
if not columns:
raise ValueError("No column references.")
return columns
@staticmethod
def _format_instruction(instruction: str, columns: List[str]) -> str:
"""Extracts column names enclosed in curly braces from the user instruction.
For example, `_format_instruction(["city", "continent"], "{city} is in {continent}")
== "city is in continent"`
"""
return instruction.format(**{col: col for col in columns})
@staticmethod
def _validate_model(model):
from bigframes.ml.llm import GeminiTextGenerator
if not isinstance(model, GeminiTextGenerator):
raise TypeError("Model is not GeminiText Generator")
@staticmethod
def _confirm_operation(row_count: int):
"""Raises OperationAbortedError when the confirmation fails"""
import bigframes # Import in the function body to avoid circular imports.
threshold = bigframes.options.compute.ai_ops_confirmation_threshold
if threshold is None or row_count <= threshold:
return
if bigframes.options.compute.ai_ops_threshold_autofail:
raise exceptions.OperationAbortedError(
f"Operation was cancelled because your work estimate is {row_count} rows, which exceeds the threshold {threshold} rows."
)
# Separate the prompt out. In IDE such VS Code, leaving prompt in the
# input function makes it less visible to the end user.
print(f"This operation will process about {row_count} rows.")
print(
"You can raise the confirmation threshold by setting `bigframes.options.compute.ai_ops_confirmation_threshold` to a higher value. To completely turn off the confirmation check, set the threshold to `None`."
)
print("Proceed? [Y/n]")
reply = input().casefold()
if reply not in {"y", "yes", ""}:
raise exceptions.OperationAbortedError("Operation was cancelled.")