apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-19 10:11:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fixed VertexAI token counts

Browse source
This commit is contained in:
JackColquitt 2024-09-29 11:17:08 -07:00
parent 24b80f4789
commit 577391183c
1 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
trustgraph/model/text_completion/vertexai/llm.py
View file

@ -142,14 +142,14 @@ class Processor(ConsumerProducer):
with __class__.text_completion_metric.time(): with __class__.text_completion_metric.time():
resp = self.llm.generate_content( response = self.llm.generate_content(
prompt, generation_config=self.generation_config, prompt, generation_config=self.generation_config,
safety_settings=self.safety_settings safety_settings=self.safety_settings
) )
resp = resp.text resp = response.text
inputtokens = resp.usage_metadata.prompt_token_count inputtokens = int(response.usage_metadata.prompt_token_count)
outputtokens = resp.usage_metadata.candidates_token_count outputtokens = int(response.usage_metadata.candidates_token_count)
print(resp, flush=True) print(resp, flush=True)
print(f"Input Tokens: {inputtokens}", flush=True) print(f"Input Tokens: {inputtokens}", flush=True)
print(f"Output Tokens: {outputtokens}", flush=True) print(f"Output Tokens: {outputtokens}", flush=True)