Last Updated on 2023-09-21 by Clay
前言
現在的生成式模型越來越厲害的，各個獨立研究人員也都部署起了一個又一個的開源大型語言模型（LLMs）。但是在使用大型語言模型做推理、生成回覆時，要是真要去等待一個比較長的輸出，那是真的挺花時間的。
實際上，像 ChatGPT 那樣的串流式（stream）輸出、一次把一段生成的 tokens 吐出，絕對是讓使用者體驗更上一層樓的好方式。
作為開源模型界的 GitHub，HuggingFace 自然注意到了這個需求。在 HuggingFace 所提供的 transformers 4.30.1 中，提供了以下兩種接口給 model.generate()：
- TextStreamer: 能夠直接在標準輸出（stdout）中印出模型生成的回覆
- TextIteratorStreamer: 能夠使用 thread 去設置生成任務，並以自定義的迭代方式（通常為 for 迴圈）來印出模型生成的回覆
對於測試模型，TextStreamer 就足夠了；對於真實上線產品，勢必得使用 TextIteratorStreamer。
使用方法
TextStreamer
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
# Tokenizer and Model
pretrained_model_name_or_path = "sshleifer/tiny-gpt2"
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path)
# Tokenized
text = "How are you?"
inputs = tokenizer(text, return_tensors="pt")
streamer = TextStreamer(tokenizer=tokenizer)
# Generation
model.generate(**inputs, streamer=streamer, max_new_tokens=50)
輸出會直接印在螢幕上。
TextIteratorStreamer
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
# Tokenizer and Model
pretrained_model_name_or_path = "sshleifer/tiny-gpt2"
tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path)
# Tokenized
text = "How are you?"
inputs = tokenizer(text, return_tensors="pt")
streamer = TextIteratorStreamer(tokenizer=tokenizer)
# Arguments
inputs.update({"streamer": streamer, "max_new_tokens": 50})
# Generation
thread = Thread(target=model.generate, kwargs=inputs)
thread.start()
for token in streamer:
print(token, end="")
thread.join()
在這裡直接印出是我選擇的表現方式，實際上可以自由選擇各種方式來表示串流資料 —— 比方說在真實上線環境中，就是讓其呈現在頁面上。