OpenAI o1を再現しよう（Reasoningモデルの作り方）

import os
from datasets import load_dataset, Dataset
from copy import deepcopy
import re
import time
from openai import OpenAI
from google.colab import userdata
import math
from tqdm.auto import tqdm
import json

# DeepInfra APIの設定
openai = OpenAI(
    api_key=userdata.get("DEEPINFRA_API_KEY"),
    base_url="https://api.deepinfra.com/v1/openai",  # DeepInfraのエンドポイント
)

# データ取得
dataset = load_dataset("nvidia/OpenMathInstruct-2", split="train")

# problem_sourceがaugmented_gsm8kでないものを抽出
filtered_dataset = dataset.filter(lambda x: x["problem_source"] != "augmented_gsm8k")

# 上位1000行のみ使用（必要に応じて変更可能）
filtered_dataset = filtered_dataset.select(range(min(1000, len(filtered_dataset))))

def translate_problem(problem_text):
    # Qwen/Qwen2.5-72B-Instructモデルを利用して翻訳
    prompt = f"以下を日本語訳してください。\n\n{problem_text}\n\n制約：\n・注釈や説明は不要です。"
    response = openai.chat.completions.create(
        model="Qwen/Qwen2.5-72B-Instruct",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=1024,
        temperature=0.7
    )
    return response.choices[0].message.content.strip()

# 新しい列problem_jaを作成
filtered_dataset = filtered_dataset.map(
    lambda x: {"problem_ja": translate_problem(x["problem"])},
    # batched=False
)

# HuggingFaceへPush
filtered_dataset.push_to_hub("OpenMathInstruct-2-CoT-JA")

# データ取得
dataset = load_dataset("HachiML/OpenMathInstruct-2-JA", split="train")

# ステップバイステップの回答を生成する関数
def generate_answer(problem_ja, generated_solution):
    # Qwen/Qwen2.5-72B-Instructを使用し、指定されたプロンプトで生成
    prompt = f"""数学の問題と模範回答が与えられます。模範回答をステップバイステップの形式に書き換えてください。

# 問題：
{problem_ja}

# 模範回答(英語)：
{generated_solution}

# 制約：
・回答は日本語で出力してください。
・推論のステップ毎に"\n\n"で区切ってください。
・途中に必ず「待って間違えたかもしれない。XXの見直しをします。」のようん見直しのステップを含めてください。
・推論毎のタイトルは不要です。
・必ず作成した回答文のみ出力してください。"""

    response = openai.chat.completions.create(
        model="Qwen/Qwen2.5-72B-Instruct",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=2048,
        temperature=0.7
    )
    return response.choices[0].message.content.strip()

# 生成された回答から答えを抽出する
def extract_boxed_answer(generated_solution):
    import re
    start_pattern = r'\\boxed\{'
    start_match = re.search(start_pattern, generated_solution)
    if not start_match:
        return None

    start_index = start_match.end()
    brace_count = 1
    i = start_index

    # 括弧の対応をカウントしながら探索
    while i < len(generated_solution) and brace_count > 0:
        if generated_solution[i] == '{':
            brace_count += 1
        elif generated_solution[i] == '}':
            brace_count -= 1
        i += 1

    if brace_count != 0:
        # 対応する閉じ括弧がない場合
        return None

    # \boxed{ ... } の中身を抽出
    content = generated_solution[start_index:i-1]

    # 不要な見た目用コマンドを削除
    # たとえば \left, \rightを除去
    # 必要に応じてパターンを追加可能
    content = re.sub(r'\\(left|right)', '', content)

    # 中の空白（スペース、タブ、改行）をすべて削除
    content = re.sub(r'\s+', '', content)

    return content.strip()

# 回答生成（失敗した場合最大3回まで再実行）
def try_generate_correct_answer(problem_ja, expected_answer, original_generated_solution):
    # 最大3回試行し、最後の(3回目)は必ず保存する
    max_retry = 3
    last_ans = None
    for i in range(max_retry):
        ans = generate_answer(problem_ja, original_generated_solution)
        boxed_val = extract_boxed_answer(ans)
        # 中の空白（スペース、タブ、改行）をすべて削除して比較
        expected_answer_stripped = None if expected_answer is None else re.sub(r"\s+", "", expected_answer)
        boxed_val_stripped = None if boxed_val is None else re.sub(r"\s+", "", boxed_val)
        if expected_answer_stripped == boxed_val_stripped:
            # 正解が見つかった場合、その時点で返す
            return ans, boxed_val, True
        last_ans = ans  # 正解でなければ次へ
        time.sleep(1)
    # 3回試行してすべて不一致の場合、最後の回答を返す（不一致だが保存）
    return last_ans, extract_boxed_answer(last_ans), False

# １行ごとの回答生成、確認のフロー
def process_row(example):
    problem_ja = example.get("problem_ja", None)
    expected_answer = example.get("expected_answer", None)
    original_generated_solution = example.get("generated_solution", None)

    if original_generated_solution is None or problem_ja is None or expected_answer is None:
        # 必要な情報がない場合はNoneを返す
        return {"generated_solution_cot": None, "cot_answer": None}

    final_ans, final_boxed, correct_flg = try_generate_correct_answer(problem_ja, expected_answer, original_generated_solution)
    # final_ans: 最後の回答(成功時は成功回答、失敗時は3回目の回答)
    # final_boxed: 最後の回答におけるboxedの中身
    return {
        "generated_solution_cot": final_ans,
        "cot_answer": final_boxed,
        "correct_flg": correct_flg
    }

# バッチごとに処理し、全データを蓄積しながらプッシュ
def process_dataset_with_batches(dataset, batch_size=100, repo_id="OpenMathInstruct-2-CoT-JA"):
    total_batches = math.ceil(len(dataset) / batch_size)
    processed_data = []

    temp_save_path = "processed_temp.json"

    # 中断再開用ローディング
    if os.path.exists(temp_save_path):
        with open(temp_save_path, "r") as f:
            processed_data = json.load(f)
        print(f"Loaded {len(processed_data)} previously processed examples.")

    start_batch = len(processed_data) // batch_size

    for batch_idx in range(start_batch, total_batches):
        start_idx = batch_idx * batch_size
        end_idx = min((batch_idx + 1) * batch_size, len(dataset))

        batch = dataset.select(range(start_idx, end_idx))
        print(f"Processing batch {batch_idx + 1}/{total_batches}...")
        processed_batch = batch.map(process_row, desc=f"Processing batch {batch_idx + 1}/{total_batches}")
        processed_data.extend(processed_batch)

        # 中間データ保存
        with open(temp_save_path, "w") as f:
            json.dump(processed_data, f)

        print(f"Pushing batch {batch_idx + 1}/{total_batches} to hub...")
        try:
            temp_dataset = Dataset.from_list(processed_data)
            temp_dataset.push_to_hub(repo_id, commit_message=f"Batch {batch_idx + 1}")
            print(f"Batch {batch_idx + 1} pushed successfully!")
        except Exception as e:
            print(f"Error pushing batch {batch_idx + 1}: {e}")
            break

    if os.path.exists(temp_save_path):
        os.remove(temp_save_path)

    print("All batches processed and pushed successfully.")
    return Dataset.from_list(processed_data)

# ステップバイステップの回答生成を実行
filtered_dataset = process_dataset_with_batches(dataset, batch_size=10)

from datasets import load_dataset, Dataset

# データセット取得
dataset = load_dataset("HachiML/OpenMathInstruct-2-CoT-JA", split="train")

import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig

# モデルとトークナイザのロード
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# 追加する特殊トークンを定義
special_tokens = ["<Thought>", "</Thought>", "<Output>", "</Output>"]
tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})

# モデルロード（サイズやGPUメモリにより適宜調整）
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
model.resize_token_embeddings(len(tokenizer))

# データセット読み込み
raw_dataset = load_dataset("HachiML/OpenMathInstruct-2-CoT-JA", split="train")

# correct_flgがTrueのものだけフィルタ
filtered_dataset = raw_dataset.filter(lambda x: x["correct_flg"] == True)

# gsm8kを除外(データ生成の際に除外し忘れていた)
filtered_dataset = filtered_dataset.filter(lambda x: x["problem_source"] != "gsm8k")

# messages列を作成
def create_messages(ex):
    prompt = ex["problem_ja"]
    generated_solution_cot = ex["generated_solution_cot"]
    cot_answer = ex["cot_answer"]
    assistant_content = f"<Thought>\n{generated_solution_cot}\n\n</Thought>\n<Output>\n{cot_answer}\n</Output>"

    messages = [
        {"role": "system", "content": "You are a helpful and harmless assistant. You should think step-by-step."},
        {"role": "user", "content": prompt},
        {"role": "assistant", "content": assistant_content}
    ]
    return {"messages": messages}

processed_dataset = filtered_dataset.map(
    create_messages,
    batched=False,
    remove_columns=filtered_dataset.column_names
)

new_model_name = "QwQ-CoT-0.5B-JA-v1.1"
train_batch_size = 4
save_steps = 80 / train_batch_size
logging_steps = 80 / train_batch_size

training_args = SFTConfig(
    output_dir="./"+new_model_name,
    overwrite_output_dir=True,
    per_device_train_batch_size=train_batch_size,
    gradient_accumulation_steps=1,
    learning_rate=1e-5,
    lr_scheduler_type="cosine",
    num_train_epochs=2,
    save_steps=save_steps,
    save_total_limit=2,
    logging_steps=logging_steps,
    bf16=True,  # BF16を使用（対応GPUが必要）
    report_to="wandb",  # wandbへログ送信
    max_seq_length=2048
)

# SFTTrainerにmessagesを含むデータセットを直接渡す
# SFTTrainerが内部でapply_chat_templateとトークナイズを行う
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=processed_dataset,
    tokenizer=tokenizer,
)

trainer.train()

# 学習完了後、モデルをBF16へ変換
model.to(torch.bfloat16)

# 学習完了後Hugging Face Hubにプッシュ
trainer.push_to_hub(new_model_name)

from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import HfApi, Repository

# モデルIDを定義
model_id = "HachiML/QwQ-CoT-0.5B-JA-v1.1"

# トークナイザーとモデルをロード
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

# ステップ区切りトークンIDの準備
step_separators = ["。\n\n", "]\n\n", "\n\n"]
step_separator_ids = [tokenizer.encode(step_separator, add_special_tokens=False)[0] for step_separator in step_separators]

# configにstep_separatorsを追加
model.config.step_separator_ids = step_separator_ids
model.config._name_or_path = model_id

# push
model.push_to_hub(model_id)

# モジュールのインポート
from reasoning_model import ReasoningModelForCausalLM
from tree_utils import print_tree_with_best_path
from transformers import AutoTokenizer

# tokenizerとmodelの準備
model_name = "HachiML/QwQ-CoT-0.5B-JA-v1.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = ReasoningModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)

system_prompt = "You are a helpful and harmless assistant. You should think step-by-step."  # 固定を推奨
prompt = "231 + 65*6 = ?"
messages = [
    {"role": "system", "content": system_prompt},
    {"role": "user", "content": prompt}
]

# chat_templateとtokenize
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# MCTSを用いて生成（Google ColabのT4インスタンスで1分程度かかる）
final_tokens, final_node = model.generate(
    **model_inputs,
    iterations_per_step=5,      # 1推論ステップの探索に何回シミュレーションを行うか。長いほど精度が高まる可能性はあるが、推論時間が伸びる。
    max_iterations=15,          # 推論ステップの上限: 0.5Bモデルの場合、そこまで長いステップの推論はできないため10~15くらいが妥当。
    mini_step_size=32,          # mini-step: 32tokens。Step as Action戦略を採用する場合、ここを512など大きな数字にする。（実行時間が伸びるため非推奨）
    expand_threshold=0,         # ノードを拡張するために必要なそのノードの訪問回数の閾値。デフォルトの0で、拡張には1回以上の訪問が求められる。基本デフォルトで良い。
    step_separator_ids=None,    # Reasoning Action StrategyでStep as Actionを採用するときの区切りとなるトークンのIDリスト。NoneでモデルConfigの値を利用するため、変更は非推奨。Step as Action不採用時には[]を設定する。
)

# 結果をテキスト化
final_text = tokenizer.decode(final_tokens, skip_special_tokens=True)
print("=== 最終生成テキスト ===")
print(final_text)

# GitHubリポジトリからクローン
!git clone https://github.com/Hajime-Y/reasoning-model.git
%cd reasoning-model

# 必要なライブラリをpipインストール
!pip install transformers torch numpy datasets

# Pythonパスにカレントディレクトリを追加
import sys
sys.path.append('.')  # reasoning_model.py と同じディレクトリにいるため

from reasoning_model import ReasoningModelForCausalLM
from tree_utils import print_tree_with_best_path
from evaluate.evaluate_mgsm import evaluate_model_on_mgsm
from transformers import AutoTokenizer

MODEL_NAME = "HachiML/QwQ-CoT-0.5B-JA-v1.1"

# 設定
iterations_per_step = 5
max_iterations = 15
mini_step_size = 32
expand_threshold = 0
step_separator_ids = None
system_prompt = "You are a helpful and harmless assistant. You should think step-by-step."

repo_id = f"mgsm_250-{MODEL_NAME.split('/')[-1]}-MCTS-ips{iterations_per_step}-mi{max_iterations}-mss{mini_step_size}-et{expand_threshold}"
if step_separator_ids is None:
    repo_id += "-sa"

# 評価の実行
results_dataset, metrics = evaluate_model_on_mgsm(
    model_name=MODEL_NAME,
    repo_id=repo_id,
    iterations_per_step=iterations_per_step,
    max_iterations=max_iterations,
    mini_step_size=mini_step_size,
    expand_threshold=expand_threshold,
    step_separator_ids=step_separator_ids,
    system_prompt=system_prompt
)