GradioとTransformersで Mistral-Small-3.1-24B-Instruct-2503 を使う

使用したPC

プロセッサ	Intel(R) Core(TM) i7-12700H
実装 RAM	32.0 GB
GPU		RTX 3080 Laptop (VRAM 16GB)

VRAM 16GBで動かすためにbitsandbytesで4bit量子化を行っています。

画像を扱うとVRAM消費量が増えるため今回はマルチモーダル非対応です。

Python環境構築

pip install torch==2.6.0+cu124 torchvision==0.21.0+cu124 --index-url https://download.pytorch.org/whl/cu124
pip install git+https://github.com/huggingface/transformers
pip install accelerate gradio bitsandbytes

accelerate==1.5.2
bitsandbytes==0.45.3
gradio==5.22.0
torch==2.6.0+cu124
torchvision==0.21.0+cu124
transformers @ git+https://github.com/huggingface/transformers@8e97b44087341837a2a9503fed9469c3abfd9eb8

Pythonスクリプト

以下のスクリプトを実行するだけです。

「tritonがない」と注意されます。
Windowsにtritonがインストールできない訳じゃないですが無視して良いと思います。

import gradio as gr
from transformers import AutoModelForImageTextToText, AutoProcessor, TextIteratorStreamer, BitsAndBytesConfig
from threading import Thread
import torch

model_name = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"

system_prompt_text = "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"
init = {
    "role": "system",
    "content": system_prompt_text
}

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

model = AutoModelForImageTextToText.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    quantization_config=quantization_config
)

tokenizer = AutoProcessor.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]
):
    text = tokenizer.apply_chat_template(
        history,
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = tokenizer(
        text=text,
        return_tensors="pt"
    ).to(0, dtype=torch.float16)

    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=2048,
        do_sample=True,
        temperature=0.15,
        pad_token_id=2, # Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.
    )
    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:
    gr.Markdown("# Mistral-Small-3.1-24B-Instruct-2503")
    chatbot = gr.Chatbot(
        type="messages",
        latex_delimiters = [
            {"left": "[", "right": "\\]", "display": True }
        ]
        
    )
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot], value="新しいチャットを開始")
    
    msg.submit(
        user, [msg, chatbot], [msg, chatbot], queue=False
    ).then(
        bot, chatbot, chatbot
    )

demo.launch()