Skip to content

Latest commit

 

History

History
40 lines (33 loc) · 1.6 KB

File metadata and controls

40 lines (33 loc) · 1.6 KB

transformers-stream-generator

PyPI - Python Version PyPI GitHub license badge Blog Downloads

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