Przeglądaj źródła

Refactor:Improve the chat stream logic for NvidiaCV (#9242)

### What problem does this PR solve?

Improve the chat stream logic for NvidiaCV

### Type of change


- [x] Refactoring
tags/v0.20.1
Stephen Hu 2 miesięcy temu
rodzic
commit
0a303d9ae1
No account linked to committer's email address
1 zmienionych plików z 4 dodań i 1 usunięć
  1. 4
    1
      rag/llm/cv_model.py

+ 4
- 1
rag/llm/cv_model.py Wyświetl plik

@@ -623,15 +623,18 @@ class NvidiaCV(Base):
return "**ERROR**: " + str(e), 0

def chat_streamly(self, system, history, gen_conf, images=[], **kwargs):
total_tokens = 0
try:
response = self._request(self._form_history(system, history, images), gen_conf)
cnt = response["choices"][0]["message"]["content"]
if "usage" in response and "total_tokens" in response["usage"]:
total_tokens += response["usage"]["total_tokens"]
for resp in cnt:
yield resp
except Exception as e:
yield "\n**ERROR**: " + str(e)

yield response["usage"]["total_tokens"]
yield total_tokens


class AnthropicCV(Base):

Ładowanie…
Anuluj
Zapisz