ご家庭用LLMでストリーミングする方法

ネットで軽く探したけど書いてなかったのでご家庭用LLMを試す時にREPLでストリーミングしたい(なぜなら長考に入ってしまうから)時に使うコード。
ELYZAのやつを例にするけど他も同じはず。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer,TextIteratorStreamer

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

model_name = "elyza/ELYZA-japanese-Llama-2-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")

if torch.cuda.is_available():
    model = model.to("cuda")


streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
pipe = pipeline(
    "text-generation",
    model=base_model,
    tokenizer=tokenizer,
    max_length=2048,
    temperature=0.6,
    pad_token_id=tokenizer.eos_token_id,
    top_p=0.95,
    repetition_penalty=1.2,
    streamer=streamer
)

from threading import Thread
def q(text):
    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
    prompt = "{bos_token}{b_inst} {system}{prompt} {e_inst} ".format(
        bos_token=tokenizer.bos_token,
        b_inst=B_INST,
        system=f"{B_SYS}{DEFAULT_SYSTEM_PROMPT}{E_SYS}",
        prompt=text,
        e_inst=E_INST,
    )
    with torch.no_grad():
        token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
        def generate(input):    
            output_ids=model.generate(input, streamer=streamer, 
                            pad_token_id=tokenizer.eos_token_id, max_length=4096, temperature=0.9, top_p=0.8,
                            eos_token_id=tokenizer.eos_token_id,
                            repetition_penalty=1.25)
        t = Thread(target=generate, args=(token_ids.to(model.device),))
        t.start()
    for new_text in streamer:
        print( new_text.replace(" ",""),end="")

q("UberEats配達員が主人公の面白い物語を考えてください")

これでこんな感じで返ってくる。

>>> q("UberEats配達員が主人公の面白い物語を考えてください")
UberEats配達員、ジョーイの話です。

ジョーイは毎日朝から晩までエッグベネディクトとコーヒーを配達している。そんな中、新しいお客さんに会うことができた。

「ハロウィンパレード」と書かれたTシャツを着た男性だ。彼は自分の住所を入力せず、ジョーイに料理を渡した。

「今度、私の家に行ってほしいものがあるんだけど…」言葉を切る前に、彼女は口を開いた。

「わぉ！めちゃくちゃ楽しみ～✨」嬉しそうに返事をしたジョーイは、隣り合うマンションへ向かった。

到着後、2人は一緒に食べ始めた。美味しいと大絶賛され、気持ちよく支払ってもらえた。最高の1日だった。

enjoy!