Skip to content

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

License

Notifications You must be signed in to change notification settings

AIxyz/transformers-stream-generator

 
 

Repository files navigation

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

About

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

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 99.8%
  • Shell 0.2%