Axcxeptが公開している「Llama-3.1-8B-EZO-1.1-it」をGradioを使ってローカルで使用する

はじめに

「Llama-3.1-8B-EZO-1.1-it」は日本語性能向上のためにLlama3.1-8Bをファインチューニングしたモデルです。

こちらになります。
huggingface.co

今のところLlama3.1の小さいモデル（8B）で日本語性能を向上させたものは少ないです。

さっそくGradioから使ってみました。

Pythonスクリプト

import gradio as gr
from transformers import pipeline, TextIteratorStreamer
from threading import Thread
import torch

system_prompt_text = "あなたは誠実で優秀な日本人のアシスタントです。特に指示が無い場合は、原則日本語で回答してください。"
init = {
    "role": "system",
    "content": system_prompt_text,
}

model_id = "AXCXEPT/Llama-3.1-8B-EZO-1.1-it"

generator = pipeline(
    "text-generation",
    model=model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16
)
streamer = TextIteratorStreamer(generator.tokenizer, skip_prompt=True, skip_special_tokens=True)

def call_llm(
    message: str,
    history: list[dict],
    max_tokens: int
):
    history_openai_format = []
    if len(history) == 0:
        history_openai_format.append(init)
        history_openai_format.append({"role": "user", "content": message})
    else:
        history_openai_format.append(init)
        for human, assistant in history:
            history_openai_format.append({"role": "user", "content": human})
            history_openai_format.append({"role": "assistant", "content": assistant})
        history_openai_format.append({"role": "user", "content": message})

    params = {
        "max_new_tokens": max_tokens,
        "streamer": streamer,
        "pad_token_id": generator.tokenizer.eos_token_id
    }

    thread = Thread(target=generator, kwargs=dict(text_inputs=history_openai_format, **params))

    thread.start()

    generated_text = ""
    for new_text in streamer:
        generated_text += new_text
        yield generated_text

demo = gr.ChatInterface(
    fn=call_llm,
    title="Axcxept Llama-3.1-8B-EZO-1.1-it",
    stop_btn="Stop Generation",
    cache_examples=False,
    additional_inputs_accordion=gr.Accordion(
        label="Parameters", open=False, render=False
    ),
    additional_inputs=[
        gr.Slider(
            minimum=1,
            maximum=4096,
            step=1,
            value=512,
            label="Max tokens",
            visible=True,
            render=False,
        )
    ],
)
demo.launch(share=False)