https://touch-sp.hatenablog.com/entry/2025/01/31/190019

はじめに

Gradioでは出力に以下のような文字が出たらうまく表示されません。

<think></think>

<answer></answer>

回避するためには一工夫が必要です。
github.com

使用したPC

Windows 11
CUDA 12.4
Python 3.12

Pythonライブラリ

accelerate==1.3.0
gradio==5.16.0
torch==2.6.0+cu124
torchao==0.8.0
transformers==4.48.3
triton @ file:///D:/triton-3.0.0-cp312-cp312-win_amd64.whl

tritonはこちらからファイルをダウンロードしてインストールしました。
こちらからTriton3.2.0をダウンロードしてインストールすると以下の警告がでました。

WARNING: Failed to find MSVC.

Pythonスクリプト

今回はサイバーエージェントが公開している「DeepSeek-R1-Distill-Qwen-14B-Japanese」をGradioで使います。

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

model_name = "DeepSeek-R1-Distill-Qwen-14B-Japanese"

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

quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    torch_dtype=torch.bfloat16,
    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=4096,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id,
    )
    history.append({"role": "assistant", "content": ""})

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

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

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

demo.launch()