gpt-oss-20b 紹介・検証：OpenAIが公開した新しいOSS LLMの実力と導入方法

uv add git+https://github.com/huggingface/transformers "git+https://github.com/triton-lang/triton@v3.4.0#subdirectory=python/triton_kernels" "triton>=3.4.0" kernels

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

LLMU 形式 (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
import logging
from pathlib import Path
from typing import List

import pandas as pd
from tqdm import tqdm
from transformers.pipelines import pipeline
from transformers.pipelines.base import Pipeline
from transformers.models.auto.tokenization_auto import AutoTokenizer
from transformers.models.auto.modeling_auto import AutoModelForCausalLM


# ---------- 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:"""

SYSTEM_PROMPT = "Reasoning: high"
# -----------------------------

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

def extract_letter(text: str) -> str:
    """gpt-oss は出力の最後に assistantfinal と出力され、その後に回答の文字が続くので、
    その文字を抽出する。大文字の A, B, C, D"""
    m = LETTER_RE.search(text)
    return m.group(1) 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)
    return pipeline(
        "text-generation",
        model=model_name,
        torch_dtype="auto",
        device_map=device,
        eos_token_id=tokenizer.eos_token_id,
        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()
    
    logging.basicConfig(
        filename='app.log',              # 出力先ファイル
        level=logging.INFO,              # 出力レベル
        format='%(asctime)s [%(levelname)s] %(message)s'  # フォーマット
    )

    # 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 = [
        [
            {
                "role": "system",
                "content": SYSTEM_PROMPT
            },
            {
                "role": "user",
                "content": 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, max_new_tokens=args.max_new_tokens)
        durations.extend([time.perf_counter() - t0] * len(batch))
        if outputs:
            for output in outputs:
                # pipeline の戻り値は list[dict] or str depending on version
                logging.info(f"Model output: {output}")
                content = str(output[0]["generated_text"] if isinstance(output, list) else output)
                predictions.append(extract_letter(content))

    # 結果集計
    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") / f"{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()