From fec9bab20041f197dc9ee04b2554cc0f99fc3d2a Mon Sep 17 00:00:00 2001 From: "Joel Z. Leibo" Date: Thu, 5 Sep 2024 01:21:25 -0700 Subject: [PATCH] Print recoverable exceptions in mistral model wrapper. PiperOrigin-RevId: 671273229 Change-Id: Ie346c3d7460d5b5fc3c8ced789ae175a5018794a --- concordia/language_model/mistral_model.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/concordia/language_model/mistral_model.py b/concordia/language_model/mistral_model.py index e9f5ca9e..8fa6af25 100644 --- a/concordia/language_model/mistral_model.py +++ b/concordia/language_model/mistral_model.py @@ -31,6 +31,7 @@ # At least one Mistral model supports completion mode. COMPLETION_MODELS = ( 'codestral-latest', + 'codestral-2405', ) @@ -64,7 +65,7 @@ def __init__( self._text_model_name = model_name self._measurements = measurements self._channel = channel - self._client = mistralai.Mistral(api_key=api_key) + self._client = mistralai.Mistral(api_key=self._api_key) self._choice_model_name = self._text_model_name if use_codestral_for_choices: @@ -109,7 +110,8 @@ def _complete_text( stop=terminators, random_seed=seed, ) - except mistralai.models.sdkerror.SDKError: + except mistralai.models.sdkerror.SDKError as err: + print(f' Exception: {err}') continue else: result = response.choices[0].message.content @@ -172,7 +174,8 @@ def _chat_text( max_tokens=max_tokens, random_seed=seed, ) - except mistralai.models.sdkerror.SDKError: + except mistralai.models.sdkerror.SDKError as err: + print(f' Exception: {err}') continue else: return response.choices[0].message.content