pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor: improve how NvidiaCV calculate res total token counts (#10455)

Browse Source 
### What problem does this PR solve?
improve how NvidiaCV calculate res total token counts

### Type of change
- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-10-10 11:03:40 +08:00
committed by GitHub
parent f631073ac2
commit 6ab4c1a6e9
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
rag/llm/cv_model.py
View File

@ -614,7 +614,7 @@ class NvidiaCV(Base):
response = response.json() response = response.json()
return ( return (
response["choices"][0]["message"]["content"].strip(), response["choices"][0]["message"]["content"].strip(),
response["usage"]["total_tokens"], total_token_count_from_response(response),
) )
def _request(self, msg, gen_conf={}): def _request(self, msg, gen_conf={}):
@ -637,7 +637,7 @@ class NvidiaCV(Base):
response = self._request(vision_prompt) response = self._request(vision_prompt)
return ( return (
response["choices"][0]["message"]["content"].strip(), response["choices"][0]["message"]["content"].strip(),
response["usage"]["total_tokens"], total_token_count_from_response(response)
) )
def chat(self, system, history, gen_conf, images=[], **kwargs): def chat(self, system, history, gen_conf, images=[], **kwargs):
@ -645,7 +645,7 @@ class NvidiaCV(Base):
response = self._request(self._form_history(system, history, images), gen_conf) response = self._request(self._form_history(system, history, images), gen_conf)
return ( return (
response["choices"][0]["message"]["content"].strip(), response["choices"][0]["message"]["content"].strip(),
response["usage"]["total_tokens"], total_token_count_from_response(response)
) )
except Exception as e: except Exception as e:
return "**ERROR**: " + str(e), 0 return "**ERROR**: " + str(e), 0
@ -656,7 +656,7 @@ class NvidiaCV(Base):
response = self._request(self._form_history(system, history, images), gen_conf) response = self._request(self._form_history(system, history, images), gen_conf)
cnt = response["choices"][0]["message"]["content"] cnt = response["choices"][0]["message"]["content"]
if "usage" in response and "total_tokens" in response["usage"]: if "usage" in response and "total_tokens" in response["usage"]:
total_tokens += response["usage"]["total_tokens"] total_tokens += total_token_count_from_response(response)
for resp in cnt: for resp in cnt:
yield resp yield resp
except Exception as e: except Exception as e: