Making TextIteratorStreamer in HuggingFace work with StreamingResponse in FastAPI

[maziyarpanahi]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google "How to X in FastAPI" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import uvicorn

from fastapi import FastAPI, Depends, HTTPException, status, Request
from fastapi.responses import StreamingResponse


from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import torch

model_id = "declare-lab/flan-alpaca-xl"
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
print("Running on device:", torch_device)
print("CPU threads:", torch.get_num_threads())


model = AutoModelForSeq2SeqLM.from_pretrained(
    model_id, load_in_8bit=True, device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_id)


def run_generation(user_text):
    # Get the model and tokenizer, and tokenize the user text.
    model_inputs = tokenizer([user_text], return_tensors="pt").to(torch_device)

    # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer
    # in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread.
    streamer = TextIteratorStreamer(
        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
    )
    generate_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=250,
        do_sample=True,
        top_p=0.95,
        temperature=float(0.8),
        top_k=1,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    # Pull the generated text from the streamer, and update the model output.
    model_output = ""
    for new_text in streamer:
        model_output += new_text
        yield new_text
    return model_output


app = FastAPI()


@app.get("/")
async def root():
    return {"message": "Welcome to FastAPI!"}


@app.post("/generate")
async def generate(query: str):
    stream_response = run_generation(query)
    return StreamingResponse(stream_response, media_type="text/plain")


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8080, log_level="debug")

Description

There is a new feature in HuggingFace called TextIteratorStreamer to stream output generated via `mode.generate() to either stdout or as an iterator. It seems to be working also in Gradio:

• reference: https://huggingface.co/docs/transformers/internal/generation_utils#transformers.TextIteratorStreamer
• gradio: https://huggingface.co/spaces/joaogante/transformers_streaming

I am trying to make this work with StreamingResponse, but I had no luck. If yield new_text to get a word-by-word output I'll see nothing until it's over. If I yield model_output then I see something before the end but I am not sure it's the right way to do it.

• new_text

Screen.Recording.2023-04-18.at.19.39.52.mov

• model_output

Screen.Recording.2023-04-18.at.19.38.21.mov

Operating System

Linux

Operating System Details

No response

FastAPI Version

0.95.0

Python Version

3.9.16

Additional Context

Apologies in advance if this seems to be something obvious

[maziyarpanahi]

yield new_text actually does work, sorry for the question.

[gesanqiu]

@maziyarpanahi How you fixed this issue? Because I didn't see the new_text response streaming, it's more like all the yield chunk is stored in somewhere and response once when the streamer raised StopIteration. Or what's the exact problem is?

[maziyarpanahi]

This could be some buffering or async issue when it waits long enough or until the end of the stream to display something. I watched this video just recently and it seems to be a good direction towards streaming/asyncio/FastAPI (the API is at the end): https://www.youtube.com/watch?v=y2cRcOPHL_U

[gesanqiu]

I agree with you that this is a buffering or async issue since I see the streaming response after I add a time.sleep(1) in my client. I will watch the video, thanks a lot.

[scutzhe]

meet same problem,how to fix quickly