transformers-stream-generator icon indicating copy to clipboard operation
transformers-stream-generator copied to clipboard

This is a text generation method which returns a generator, streaming out each token in real-time during inference, based on Huggingface/Transformers.

transformers-stream-generator

PyPI - Python Version PyPI GitHub license badge Blog

Description

This is a text generation method which returns a generator, streaming out each token in real-time during inference, based on Huggingface/Transformers.

Web Demo

  • original
  • stream

Installation

pip install transformers-stream-generator

Usage

  1. just add two lines of code before your original code
from transformers_stream_generator import init_stream_support
init_stream_support()
  1. add do_stream=True in model.generate function and keep do_sample=True, then you can get a generator
generator = model.generate(input_ids, do_stream=True, do_sample=True)
for token in generator:
    word = tokenizer.decode(token)
    print(word)

Example

  • run python script example by gpt2
  • run web example by gpt2 and test in client example