feat: add client side retry to GeminiTextGenerator#1242
Conversation
bigframes/ml/llm.py
Outdated
| top_k: int = 40, | ||
| top_p: float = 1.0, | ||
| ground_with_google_search: bool = False, | ||
| retry: int = 0, |
There was a problem hiding this comment.
Can we rename it to max_retries?
|
|
||
| if (df[_ML_GENERATE_TEXT_STATUS] != "").any(): | ||
| df_succ = df[df[_ML_GENERATE_TEXT_STATUS].str.len() == 0] | ||
| df_fail = df[df[_ML_GENERATE_TEXT_STATUS].str.len() > 0] |
There was a problem hiding this comment.
Can we just use df_succ to tell whether there is any failed row? Then it can reduce one filter op.
There was a problem hiding this comment.
not really. We need df_fail for next retry, and append df_succ to df_result of current round. And I don't think DF subtract will be more effecient.
| The default is `False`. | ||
|
|
||
| max_retries (int, default 0): | ||
| Max number of retry rounds if any rows failed in the prediction. Each round need to make progress (has succeeded rows) to continue the next retry round. |
There was a problem hiding this comment.
[nit] suggest rewording
| Max number of retry rounds if any rows failed in the prediction. Each round need to make progress (has succeeded rows) to continue the next retry round. | |
| Max number of retries if the prediction for any rows failed. Each try needs to make progress (i.e. has successfully predicted rows) to continue the retry. |
|
|
||
| max_retries (int, default 0): | ||
| Max number of retry rounds if any rows failed in the prediction. Each round need to make progress (has succeeded rows) to continue the next retry round. | ||
| Each round will append newly succeeded rows. When the max retry rounds is reached, the remaining failed rows will be appended to the end of the result. |
There was a problem hiding this comment.
[nit] suggest rewording
| Each round will append newly succeeded rows. When the max retry rounds is reached, the remaining failed rows will be appended to the end of the result. | |
| Each retry will append newly succeeded rows. When the max retries are reached, the remaining rows (the ones without successful predictions) will be appended to the end of the result. |
| df = self._bqml_model.generate_text(df_fail, options) | ||
|
|
||
| if (df[_ML_GENERATE_TEXT_STATUS] != "").any(): | ||
| df_succ = df[df[_ML_GENERATE_TEXT_STATUS].str.len() == 0] |
There was a problem hiding this comment.
for readability
success = df[_ML_GENERATE_TEXT_STATUS].str.len() == 0
df_succ = df[success]
df_fail = df[~success]| break | ||
|
|
||
| df_result = ( | ||
| bpd.concat([df_result, df_succ]) if not df_result.empty else df_succ |
There was a problem hiding this comment.
Isn't the if-else redundant here? bpd.concat() should handle if any of the input dfs is empty.
There was a problem hiding this comment.
It gave me error in tests, sth related to multi-index
| df_fail = df[df[_ML_GENERATE_TEXT_STATUS].str.len() > 0] | ||
|
|
||
| if df_succ.empty: | ||
| warnings.warn("Can't make any progress, stop retrying.", RuntimeWarning) |
There was a problem hiding this comment.
Looks like this warning will surface even if the user didn't wish for any retries (the default behavior), something which is not desirable
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
b/385221339