8000 perf: eliminate count queries in llm retry by GarrettWu · Pull Request #1489 · googleapis/python-bigquery-dataframes · GitHub
[go: up one dir, main page]
Skip to content

perf: eliminate count queries in llm retry #1489

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
GarrettWu merged 2 commits into from
Mar 13, 2025
Merged

perf: eliminate count queries in llm retry #1489

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
53c51f9
performance: eliminate count queries in llm retry
GarrettWu Mar 12, 2025
262dff8
fix tests
GarrettWu Mar 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 13 additions & 22 deletions bigframes/ml/base.py
View file
Open in desktop
60E6
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
"""

import abc
from typing import Callable, cast, Mapping, Optional, TypeVar
from typing import Callable, cast, Mapping, Optional, TypeVar, Union
import warnings

import bigframes_vendored.sklearn.base
Expand Down Expand Up @@ -259,38 +259,29 @@ def _predict_and_retry(
) -> bpd.DataFrame:
assert self._bqml_model is not None

df_result = bpd.DataFrame(session=self._bqml_model.session) # placeholder
df_fail = X
for _ in range(max_retries + 1):
df_result: Union[bpd.DataFrame, None] = None # placeholder
df_succ = df_fail = X
for i in range(max_retries + 1):
if i > 0 and df_fail.empty:
break
if i > 0 and df_succ.empty:
msg = bfe.format_message("Can't make any progress, stop retrying.")
warnings.warn(msg, category=RuntimeWarning)
break

df = self._predict_func(df_fail, options)

success = df[self._status_col].str.len() == 0
df_succ = df[success]
df_fail = df[~success]

if df_succ.empty:
if max_retries > 0:
msg = bfe.format_message("Can't make any progress, stop retrying.")
warnings.warn(msg, category=RuntimeWarning)
break

df_result = (
bpd.concat([df_result, df_succ]) if not df_result.empty else df_succ
)

if df_fail.empty:
break

if not df_fail.empty:
msg = bfe.format_message(
f"Some predictions failed. Check column {self._status_col} for detailed "
"status. You may want to filter the failed rows and retry."
bpd.concat([df_result, df_succ]) if df_result is not None else df_succ
)
warnings.warn(msg, category=RuntimeWarning)

df_result = cast(
bpd.DataFrame,
bpd.concat([df_result, df_fail]) if not df_result.empty else df_fail,
bpd.concat([df_result, df_fail]) if df_result is not None else df_fail,
)
return df_result

Expand Down
0