RakutenAI-7BとGradioでチャットボットをサクッと作成｜新しい大規模言語モデル（LLM）を試してみる

!pip install -q accelerate
!pip install gradio

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer


MODEL_NAME = "Rakuten/RakutenAI-7B-chat"

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, 
                                             device_map="auto", 
                                             torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

model.eval()

requests = [
    "「馬が合う」はどう言う意味ですか",
    "How to make an authentic Spanish Omelette?",
]

system_message = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {user_input} ASSISTANT:"

for req in requests:
    input_req = system_message.format(user_input=req)
    input_ids = tokenizer.encode(input_req, return_tensors="pt").to(device=model.device)
    tokens = model.generate(
        input_ids,
        max_new_tokens=1024,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
        streamer = streamer
    )
    out = tokenizer.decode(tokens[0][len(input_ids[0]):], skip_special_tokens=True)
    print("USER:\n" + req)
    print("ASSISTANT:\n" + out)
    print()
    print()

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer


MODEL_NAME = "Rakuten/RakutenAI-7B-chat"

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, 
                                             device_map="auto", 
                                             torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

model.eval()

SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。\n"

def compose_prompt(message:str, history:list[list[str]]) -> str:
    prompt = f"{SYSTEM_PROMPT}"
    if len(history) == 0:
        prompt += f"USER:{message}\nASSISTANT:"
        return prompt
    else:
        for pair in history:
            [user, assistant] = pair
            prompt += f"USER: {user} \nASSISTANT:{assistant} <|endoftext|>\n"
        
        prompt += f"USER:{message}\nASSISTANT:"
        return prompt

  
# プロンプトを確認
prompt = compose_prompt("日本で一番高い山は？", [["一番長い川は？", "信濃川です"]])
print(prompt)

from threading import Thread

async def gen_stream(prompt: str) -> TextIteratorStreamer:
    input_ids = tokenizer(prompt, return_tensors="pt").to(device=model.device)

    # print(prompt, input_ids)
    config = dict(
        **input_ids,
        max_new_tokens=512,
        streamer=streamer,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.5,
    )

    thread = Thread(target=model.generate, kwargs=config)
    thread.start()

    return streamer

import gradio as gr 

async def chat(message: str, history: list[list[str]]):
    prompt = compose_prompt(message, history)

    st = await gen_stream(prompt)

    total_response = ""

    for output in st:
        if not output:
            continue

        total_response += output

        yield total_response

demo = gr.ChatInterface(chat).queue()
demo.launch(share=True)