arcstep / langchain_zhipuai

为了兼容Langchain,根据zhipu官方HTTP的API重新实现的Python SDK
50 stars 7 forks source link

请问大佬这个流式传输怎么通过langchain获得usage中耗费的token数啊 #6

Closed chivalry1314 closed 7 months ago

chivalry1314 commented 7 months ago

我又来了,请问大佬这个流式传输怎么通过langchain获得usage中耗费的token数啊

arcstep commented 7 months ago

感谢您的反馈!

这里确实有一点问题,invoke方法可以获得token统计,而stream方法没有获得。

我要研究一下langchain的代码回调,找出是哪里出了问题。

arcstep commented 7 months ago

我已经将写好的代码发布到最新版本 v4.1.0,你可以更新 langchain_zhipu 来体验这个功能。

以下是我编写的一个示例,你也可以直接访问 notes/tokens.ipynb 来查看相关内容。

from langchain_zhipu import ChatZhipuAI
from langchain_core.callbacks import BaseCallbackHandler

class CountTokenCallbackHandler(BaseCallbackHandler):
    def on_llm_end(self, response, **kwargs):
        print("-"*20, "Token统计", "-"*20)
        print(response)
        print("-"*20, "--------", "-"*20)

from langchain_zhipu import ChatZhipuAI
llm_glm = ChatZhipuAI(callbacks=[CountTokenCallbackHandler()])

for s in llm_glm.stream([{"role": "user", "content": "你是什么模型?"}]):
  print(s)
content='我' id='8563693888117904819'
content='使用的' id='8563693888117904819'
content='模型' id='8563693888117904819'
content='是' id='8563693888117904819'
content='清华大学' id='8563693888117904819'
content=' K' id='8563693888117904819'
content='EG' id='8563693888117904819'
content=' 实' id='8563693888117904819'
content='验' id='8563693888117904819'
content='室' id='8563693888117904819'
content='和' id='8563693888117904819'
content='智' id='8563693888117904819'
content='谱' id='8563693888117904819'
content='AI' id='8563693888117904819'
content='共同' id='8563693888117904819'
content='训练' id='8563693888117904819'
content='的' id='8563693888117904819'
content=' GL' id='8563693888117904819'
content='M' id='8563693888117904819'
content=' 模' id='8563693888117904819'
content='型' id='8563693888117904819'
content=',' id='8563693888117904819'
content='一种' id='8563693888117904819'
content='基于' id='8563693888117904819'
content=' Transformer' id='8563693888117904819'
content=' 的' id='8563693888117904819'
content='通用' id='8563693888117904819'
content='预' id='8563693888117904819'
content='训练' id='8563693888117904819'
content='语言' id='8563693888117904819'
content='模型' id='8563693888117904819'
content='。' id='8563693888117904819'
content='Transformer' id='8563693888117904819'
content=' 模' id='8563693888117904819'
content='型' id='8563693888117904819'
content='是一种' id='8563693888117904819'
content='基于' id='8563693888117904819'
content='自' id='8563693888117904819'
content='注意力' id='8563693888117904819'
content='机制的' id='8563693888117904819'
content='深度' id='8563693888117904819'
content='神经网络' id='8563693888117904819'
content='模型' id='8563693888117904819'
content=',' id='8563693888117904819'
content='经常' id='8563693888117904819'
content='用于' id='8563693888117904819'
content='处理' id='8563693888117904819'
content='序列' id='8563693888117904819'
content='数据' id='8563693888117904819'
content='。\n\n我' id='8563693888117904819'
content='可能' id='8563693888117904819'
content='用到' id='8563693888117904819'
content='最大的' id='8563693888117904819'
content='模型' id='8563693888117904819'
content='是' id='8563693888117904819'
content=' GL' id='8563693888117904819'
content='M' id='8563693888117904819'
content='-' id='8563693888117904819'
content='130' id='8563693888117904819'
content='B' id='8563693888117904819'
content=',' id='8563693888117904819'
content='具有' id='8563693888117904819'
content=' ' id='8563693888117904819'
content='130' id='8563693888117904819'
content='0' id='8563693888117904819'
content=' 亿' id='8563693888117904819'
content='参数' id='8563693888117904819'
content=',' id='8563693888117904819'
content='支持' id='8563693888117904819'
content='中' id='8563693888117904819'
content='英' id='8563693888117904819'
content='双语' id='8563693888117904819'
content='。' id='8563693888117904819'
content='我' id='8563693888117904819'
content='具体' id='8563693888117904819'
content='使用的' id='8563693888117904819'
content='模型' id='8563693888117904819'
content='规模' id='8563693888117904819'
content='视' id='8563693888117904819'
content='应用' id='8563693888117904819'
content='场景' id='8563693888117904819'
content='可能会有' id='8563693888117904819'
content='所' id='8563693888117904819'
content='变化' id='8563693888117904819'
content='。' id='8563693888117904819'
content='' id='8563693888117904819'
-------------------- Token统计 --------------------
generations=[[ChatGenerationChunk(text='我使用的模型是清华大学 KEG 实验室和智谱AI共同训练的 GLM 模型,一种基于 Transformer 的通用预训练语言模型。Transformer 模型是一种基于自注意力机制的深度神经网络模型,经常用于处理序列数据。\n\n我可能用到最大的模型是 GLM-130B,具有 1300 亿参数,支持中英双语。我具体使用的模型规模视应用场景可能会有所变化。', generation_info={'model': 'glm-4', 'created': 1712936199, 'index': 0, 'finish_reason': 'stop', 'usage': {'prompt_tokens': 9, 'completion_tokens': 91, 'total_tokens': 100}}, message=AIMessageChunk(content='我使用的模型是清华大学 KEG 实验室和智谱AI共同训练的 GLM 模型,一种基于 Transformer 的通用预训练语言模型。Transformer 模型是一种基于自注意力机制的深度神经网络模型,经常用于处理序列数据。\n\n我可能用到最大的模型是 GLM-130B,具有 1300 亿参数,支持中英双语。我具体使用的模型规模视应用场景可能会有所变化。'))]] llm_output=None run=None
-------------------- -------- --------------------

PS:

我已在 _stream 方法中增加了token统计,但并未将其加入到 llm_output 字段中,而是放到了 ChatGenerationChunk 中。若要加入 llm_output,则需要修改 stream 方法,但这将带来较大的侵入性。

由于 invoke 方法是从 llm_output 中获取token统计,因此 stream 方法和 invoke 方法在表现上可能会有所不同。

实际上,默认情况下无法看到 stream 方法中的token统计,这是因为 langchain 的 stream 方法并未提供这种功能。因此,OpenAI等其他模型的 stream 方法同样无法显示这个统计。基于这个原因,我不想直接修改 stream 方法,因为这样做会带来较大的侵入性,且在未来 langchain 更新 stream 方法时可能不易兼容。

arcstep commented 7 months ago

若想深入探究stream与invoke方法在处理llm_output方面的具体差异,可以查阅langchain的源代码,相关部分的代码非常简短: