diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index 5f487576e..328591b6d 100644 --- a/rag/llm/cv_model.py +++ b/rag/llm/cv_model.py @@ -522,6 +522,7 @@ class GeminiCV(Base): def chat_streamly(self, system, history, gen_conf, images=[]): from transformers import GenerationConfig ans = "" + response = None try: response = self.model.generate_content( self._form_history(system, history, images), @@ -537,8 +538,10 @@ class GeminiCV(Base): except Exception as e: yield ans + "\n**ERROR**: " + str(e) - yield response._chunks[-1].usage_metadata.total_token_count - + if response and hasattr(response, "usage_metadata") and hasattr(response.usage_metadata, "total_token_count"): + yield response.usage_metadata.total_token_count + else: + yield 0 class NvidiaCV(Base): _FACTORY_NAME = "NVIDIA"