オンプレLLM導入の現実解：GPUゼロでも動く日本語LLMをJMMLUで評価

You are an expert exam-taker.
Read the multiple-choice question and reply **ONLY** with the SINGLE letter
(A, B, C, or D) that correctly completes it.

Question:
{question}

A. {A}
B. {B}
C. {C}
D. {D}

Answer:

#!/usr/bin/env python3
"""
main.py

MMLU 形式 (question, A, B, C, D, answer) の CSV を読み込み、
Hugging Face Transformers の pipeline でゼロショット回答を行い精度を計測する。

CSV 必須列: question, A, B, C, D, answer
"""

from __future__ import annotations
import argparse
import re
import time
from pathlib import Path
from typing import List

import pandas as pd
from tqdm import tqdm
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    pipeline,
    Pipeline,
)


# ---------- Prompt ----------
PROMPT_TEMPLATE = """You are an expert exam-taker.
Read the multiple-choice question and reply **ONLY** with the SINGLE letter
(A, B, C, or D) that correctly completes it.

Question:
{question}

A. {A}
B. {B}
C. {C}
D. {D}

Answer:"""
# -----------------------------

LETTER_RE = re.compile(r"[A-D]")

def extract_letter(text: str) -> str:
    """LLM 出力から最初に現れる A–D を抽出"""
    m = LETTER_RE.search(text.upper())
    return m.group(0) if m else ""


def build_pipeline(
    model_name: str,
    device: str | int | None,
    max_new_tokens: int,
    **model_kwargs,
) -> Pipeline:
    """
    HuggingFace からモデルとトークナイザをロードして text-generation pipeline を生成
    device: "cpu", "cuda", "mps", or int (GPU ID). None なら auto
    model_kwargs: torch_dtype, trust_remote_code など自由に渡せる
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, **model_kwargs)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto" if device is None else None,
        **model_kwargs,
    )
    model.generation_config.pad_token_id = tokenizer.eos_token_id  # パディングトークンを設定
    return pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        device=device,
        max_new_tokens=max_new_tokens,
        do_sample=False,      # 複数の回答候補を生成しない
        temperature=0.0,
        return_full_text=False,
    )


def batch_iter(lst: List[str], n: int):
    """n 個ずつのバッチイテレータ"""
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("csv_path", help="問題 CSV パス")
    parser.add_argument(
        "--model",
        required=True,
        help="Hugging Face Model ID or local path (例: meta-llama/Meta-Llama-3-8B-Instruct)",
    )
    parser.add_argument(
        "--device",
        default=None,
        help='"cpu", "cuda", "mps" もしくは GPU ID。未指定なら自動',
    )
    parser.add_argument("--batch-size", type=int, default=4, help="推論バッチサイズ")
    parser.add_argument("--max-new-tokens", type=int, default=1, help="生成トークン数上限")
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="モデルのカスタムコードを信頼する (必要な場合のみ)",
    )
    args = parser.parse_args()

    # CSV 読み込み & 検証
    df = pd.read_csv(args.csv_path)
    required_cols = {"question", "A", "B", "C", "D", "answer"}
    if not required_cols.issubset(df.columns):
        raise ValueError(f"CSV must contain columns: {required_cols}")

    # モデルロード
    print("Loading model …")
    pipe = build_pipeline(
        model_name=args.model,
        device=args.device,
        max_new_tokens=args.max_new_tokens,
        trust_remote_code=args.trust_remote_code,
    )

    prompts = [
        PROMPT_TEMPLATE.format(
            question=row["question"],
            A=row["A"],
            B=row["B"],
            C=row["C"],
            D=row["D"],
        )
        for _, row in df.iterrows()
    ]

    predictions, durations = [], []
    print("Running inference …")
    for batch in tqdm(batch_iter(prompts, args.batch_size), total=(len(prompts) + args.batch_size - 1) // args.batch_size):
        t0 = time.perf_counter()
        outputs = pipe(batch)
        durations.extend([time.perf_counter() - t0] * len(batch))
        for out in outputs:
            # pipeline の戻り値は list[dict] or str depending on version
            text = out[0]["generated_text"] if isinstance(out, list) else out
            predictions.append(extract_letter(text))

    # 結果集計
    df["prediction"] = predictions
    df["correct"] = df["prediction"].str.upper() == df["answer"].str.upper()
    df["seconds"] = durations

    accuracy = df["correct"].mean()
    print(f"\nAccuracy: {accuracy:.2%} ({df['correct'].sum()}/{len(df)})")
    print(f"Avg latency: {df['seconds'].mean():.2f} s / question (batched)")

    model_name = args.model.split("/")[-1]  # 最後の部分をモデル名とする
    out_path = Path("./results/").with_stem(Path(args.csv_path).stem + "_results-" + model_name)
    df.to_csv(out_path.with_suffix(".csv"), index=False)
    print(f"Detailed results saved to {out_path}.csv")


if __name__ == "__main__":
    main()