elixir-nx / bumblebee

Pre-trained Neural Network models in Axon (+ 🤗 Models integration)
Apache License 2.0
1.32k stars 95 forks source link

Add token summary to text generation output #336

Closed jonatanklosko closed 7 months ago

jonatanklosko commented 7 months ago

Closes #287.

The summary looks like this token_summary: %{input: 5, output: 20, padding: 95}. When streaming, there is a new option :stream_done, which enables a final message and includes the usual result except for text, for example {:done, %{token_summary: ...}}.