「 Llama-3.1-Swallow-8B-Instruct-v0.3」をGradioを使ってローカルで使用する

関連記事

こちらは「 Llama-3.1-Swallow-8B-Instruct-v0.2」を実行した時の記事です。
touch-sp.hatenablog.com

Python環境

accelerate==1.5.2
bitsandbytes==0.45.3
gradio==5.21.0
torch==2.6.0+cu124
transformers==4.49.0

Pythonスクリプト

前回からの改善点として今回はGradioの「ClearButton」を使いました。

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
from threading import Thread
import torch

model_name = "tokyotech-llm/Llama-3.1-Swallow-8B-Instruct-v0.3"

system_prompt_text = "あなたは誠実で優秀な日本人のアシスタントです。"
init = {
    "role": "system",
    "content": system_prompt_text
}

quantization_config = BitsAndBytesConfig(load_in_8bit=True)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    quantization_config=quantization_config,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextIteratorStreamer(
    tokenizer,
    skip_prompt=True,
    skip_special_tokens=True
)

def user(
    message: str,
    history: list[dict]
):
    if len(history) == 0:
        history.insert(0, init)
    history.append(
        {
            "role": "user", 
            "content": message
        }
    )
    return "", history

def bot(
    history: list[dict]
):
    input_tensors = tokenizer.apply_chat_template(
        history,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True
    ).to(model.device)

    generation_kwargs = dict(
        **input_tensors,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.6,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id,
    )
    history.append({"role": "assistant", "content": ""})

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

    for new_text in streamer:
        history[-1]["content"] += new_text
        yield history

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(type="messages")
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot], value="新しいチャットを開始")
    
    msg.submit(
        user, [msg, chatbot], [msg, chatbot], queue=False
    ).then(
        bot, chatbot, chatbot
    )

demo.launch()

補足

LATEX表示が上手くいかない時がありました。以下のように修正しました。

chatbot = gr.Chatbot(
    type="messages",
    latex_delimiters = [{"left": "$", "right": "$", "display": True }]
)