@Blog

import os
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType

# GPUの利用可能性を確認
print("GPUが利用可能か確認中...")
if not torch.cuda.is_available():
    print("GPUが見つかりません。Fine-tuningにはGPUが強く推奨されます。")
    # GPUがない場合は、ここでスクリプトを終了するか、CPUモードで続行するか選択できます。
    # exit("GPUがないため終了します。")
else:
    print(f"GPUが利用可能です: {torch.cuda.get_device_name(0)}")

# --- 1. モデルとトークナイザーのロード ---
# Llama 3モデルのパスを指定します。Hugging Faceのモデル名（例: "meta-llama/Llama-3-8B"）
# またはローカルにダウンロードしたモデルのパスを指定してください。
MODEL_NAME = "meta-llama/Llama-3-8B" # ユーザーが入力したパスがここに挿入されます

print(f"モデルとトークナイザーをロード中: {MODEL_NAME}")

# 4bit量子化設定 (GPUメモリの節約に役立ちます)
# bnb_4bit_compute_dtypeは、Ampere以降のNVIDIA GPUに推奨されるbfloat16を使用しています。
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4", # NF4 (NormalFloat4) 量子化タイプ
    bnb_4bit_compute_dtype=torch.bfloat16 
)

# トークナイザーをロード
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
# Llama 3はデフォルトでbos_tokenを付与しないことがあるため、明示的に追加。
# また、padding_side='right'はLlamaモデルに推奨される設定です。
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# モデルをロードし、量子化設定を適用し、自動的にGPUにマッピングします。
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto", # 利用可能なデバイス（GPU）に自動的にモデルを分散
    trust_remote_code=True # リモートコードの実行を許可
)
print("モデルロード完了。")

# k-bit学習用にモデルを準備 (PEFTライブラリのため)
# gradient_checkpointingを有効にすることで、メモリ使用量をさらに削減できます。
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
print("k-bit学習用にモデルを準備しました。")

# --- 2. データセットの準備 ---
# あなたのテキスト記事ファイルが格納されているディレクトリを指定します。
# 例: 'your_article_data/' の中に 'article1.txt', 'article2.txt', ... と置かれている場合
DATA_DIR = "./your_article_data/" # ユーザーが入力したパスがここに挿入されます

print(f"データセットをロード中: {DATA_DIR}")

# 'text'形式でデータセットをロードします。指定されたディレクトリ内のすべての.txtファイルを読み込みます。
# 各ファイルが1つのエントリとして扱われます。
try:
    dataset = load_dataset('text', data_files={'train': os.path.join(DATA_DIR, '*.txt')})
    print(f"データセットのサンプル数: {len(dataset['train'])}")
except Exception as e:
    print(f"データセットのロード中にエラーが発生しました。ディレクトリとファイル形式を確認してください: {e}")
    exit("データセットロード失敗。")

# データセットをトークン化する関数
# 長い記事をモデルの最大入力長に分割します。
def tokenize_function(examples):
    # Llama 3の最大入力長は通常8192ですが、お使いのGPUのVRAMに合わせて調整してください。
    # ここでは一般的な値として2048を設定しています。
    max_length = 2048 
    # truncate=Trueで最大長を超えるテキストを切り捨てます。
    return tokenizer(examples["text"], truncation=True, max_length=max_length)

# データセットをトークン化します。
# num_procはCPUコア数に応じて並列処理を行い、処理を高速化します。
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=os.cpu_count(),
    remove_columns=["text"] # 元のテキスト列は学習に不要になるため削除します。
)
print("データセットのトークン化が完了しました。")

# --- 3. PEFT (LoRA) の設定 ---
# LoRA (Low-Rank Adaptation) は、元のモデルの重みをフリーズし、
# 小さなアダプター層を追加して学習させることで、効率的にファインチューニングを行います。
# これにより、GPUメモリの使用量を抑えつつ、高い性能を実現できます。
lora_config = LoraConfig(
    r=16, # LoRAのランク。値を大きくすると表現力が増すが、メモリ消費も増える。
    lora_alpha=32, # LoRAのスケーリング係数。rの2倍程度が推奨されることが多いです。
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], # LoRAを適用する層。Llamaモデルで一般的な層。
    bias="none", # バイアスを学習しない設定。
    lora_dropout=0.05, # ドロップアウト率。過学習を防ぐために設定します。
    task_type=TaskType.CAUSAL_LM, # タスクタイプを因果言語モデルに設定。
)

# モデルにLoRAアダプターを追加します。
model = get_peft_model(model, lora_config)
print("モデルにLoRAアダプターを適用しました。")
model.print_trainable_parameters() # 学習可能なパラメータ数を確認します。

# --- 4. 学習の実行 ---
# 学習済みモデルを保存するディレクトリ
OUTPUT_DIR = "./llama3_finetuned_model/" # ユーザーが入力したパスがここに挿入されます

# 学習の設定
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=3, # エポック数。データセットのサイズと希望する精度に応じて調整してください。
    per_device_train_batch_size=1, # GPUあたりのバッチサイズ。VRAMが少ない場合は1に設定。
    gradient_accumulation_steps=4, # 勾配を蓄積するステップ数。実質的なバッチサイズは per_device_train_batch_size * gradient_accumulation_steps になります。
    optim="paged_adamw_8bit", # 8bit AdamWオプティマイザを使用し、メモリ効率を向上させます。
    save_steps=500, # 500ステップごとにモデルを保存します。
    logging_steps=100, # 100ステップごとにログを出力します。
    learning_rate=2e-4, # 学習率。
    fp16=True, # 混合精度学習を有効化 (GPUが対応している場合)。VRAM削減と高速化に寄与します。
    max_steps=-1, # num_train_epochsに基づいて学習します。
    group_by_length=True, # 同じ長さのシーケンスをグループ化し、パディングを削減します。
    lr_scheduler_type="cosine", # 学習率スケジューラーのタイプ。
    warmup_ratio=0.03, # ウォームアップ比率。
    report_to="none", # レポート先を指定しない (wandbなどを使用しない場合)。
)

# トレーナーの初期化
# data_collatorは、モデルの入力形式に合わせてデータを整形します。
trainer = Trainer(
    model=model,
    train_dataset=tokenized_dataset["train"],
    args=training_args,
    data_collator=lambda data: {
        'input_ids': torch.stack([f['input_ids'] for f in data]),
        'attention_mask': torch.stack([f['attention_mask'] for f in data]),
        'labels': torch.stack([f['input_ids'] for f in data]), # 因果言語モデルでは、入力自体がラベルとなります。
    },
)

# 学習の開始
print("Fine-tuningを開始します...")
trainer.train()
print("Fine-tuningが完了しました。")

# --- 5. 学習済みモデルの保存 ---
# LoRAアダプターのみを保存します。これにより、ファイルサイズが小さく、効率的に管理できます。
trainer.save_model(OUTPUT_DIR)
print(f"学習済みLoRAアダプターが '{OUTPUT_DIR}' に保存されました。")

# 保存したアダプターを使って推論を行う方法の例 (コメントアウトされています):
# このコードは、ファインチューニング後にモデルをロードして推論を行うための参考例です。
# from peft import PeftModel
#
# # 元のモデルをロード (学習時と同じ量子化設定を使用します)
# base_model = AutoModelForCausalLM.from_pretrained(
#     MODEL_NAME,
#     quantization_config=bnb_config,
#     device_map="auto",
#     trust_remote_code=True
# )
#
# # 保存したLoRAアダプターを元のモデルに結合します。
# peft_model = PeftModel.from_pretrained(base_model, OUTPUT_DIR)
#
# # 推論モードに設定します。
# peft_model.eval()
#
# # テキスト生成の例
# prompt = "ローカルPCでのLlama 3ファインチューニングの利点とは"
# inputs = tokenizer(prompt, return_tensors="pt").to("cuda") # 入力をGPUに移動
#
# with torch.no_grad(): # 勾配計算を無効化し、メモリ使用量を削減
#     outputs = peft_model.generate(
#         **inputs,
#         max_new_tokens=200, # 生成する新しいトークンの最大数
#         do_sample=True, # サンプリングによる生成を有効化
#         top_p=0.9, # Nucleusサンプリングの閾値
#         temperature=0.7, # 生成の多様性を制御する温度
#         eos_token_id=tokenizer.eos_token_id # 終了トークンID
#     )
# print("\n--- 生成されたテキスト ---")
# print(tokenizer.decode(outputs[0], skip_special_tokens=True))