webbigdata/

nanochat-jp_base

# test_pretrained_jp.py

import os
import sys
sys.path.append(os.getcwd())

import torch
from nanochat.common import get_base_dir
from nanochat.checkpoint_manager import load_model_from_dir

# --- 設定 ---
MODEL_DIR_NAME = "base_checkpoints_jp"
MODEL_TAG = "d20"
STEP = None
MAX_NEW_TOKENS = 100
TEMPERATURE = 0.7
TOP_K = 50

# --- メイン実行部 ---
if __name__ == "__main__":
    print("--- 事前学習済み日本語モデル テストスクリプト ---")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用デバイス: {device}")

    use_bf16 = (device == 'cuda' and torch.cuda.is_bf16_supported())
    autocast_ctx = torch.amp.autocast(device_type=device, dtype=torch.bfloat16, enabled=use_bf16)
    if use_bf16:
        print("bfloat16がサポートされています。混合精度で推論を実行します。")

    base_dir = get_base_dir()
    checkpoints_dir = os.path.join(base_dir, MODEL_DIR_NAME)
    
    print(f"モデルを次のパスから読み込みます: {os.path.join(checkpoints_dir, MODEL_TAG)}")
    if not os.path.exists(os.path.join(checkpoints_dir, MODEL_TAG)):
        print("\nFATAL: モデルディレクトリが見つかりません。")
        print(f"ローカルの '{checkpoints_dir}' 以下に 'd20' などのディレクトリとしてモデルファイルが配置されているか確認してください。")
        sys.exit(1)

    model, tokenizer, meta = load_model_from_dir(
        checkpoints_dir, device, phase="eval", model_tag=MODEL_TAG, step=STEP
    )
    
    print("\nモデルとトークナイザーのロードが完了しました。")
    print("事前学習済みモデルは、対話ではなく『文章の続き』を生成します。")
    
    while True:
        try:
            user_input = input("\nプロンプトを入力してください (終了するにはCtrl+C): ")
            if not user_input:
                continue

            prompt_tokens = tokenizer.encode(user_input, prepend=tokenizer.get_bos_token_id())
            
            print("-" * 30)
            print("生成開始...")
            print("入力プロンプト: ", user_input, end="")

            with torch.no_grad():
                with autocast_ctx:
                    stream = model.generate(prompt_tokens, max_tokens=MAX_NEW_TOKENS, temperature=TEMPERATURE, top_k=TOP_K)
                    
                    for token in stream:
                        print(tokenizer.decode([token]), end="", flush=True)

            print("\n" + "-" * 30)

        except KeyboardInterrupt:
            print("\n終了します。")
            break
        except Exception as e:
            print(f"\nエラーが発生しました: {e}")
            break