train.py (4600B)
1 import torch 2 from datasets import load_dataset 3 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training 4 from transformers import ( 5 AutoModelForCausalLM, 6 AutoTokenizer, 7 BitsAndBytesConfig, 8 Trainer, 9 TrainingArguments, 10 ) 11 12 # --- 設定 --- 13 MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2" # 高性能なベースモデル 14 DATASET_PATH = "typst_feedback_dataset.jsonl" # 統合した学習データセット 15 OUTPUT_DIR = "./mistral_typst_adapter" # 学習結果（LoRAアダプター）の保存先 16 LORA_R = 16 # LoRAのランク（精度とメモリのトレードオフ） 17 LORA_ALPHA = 32 18 LORA_DROPOUT = 0.05 19 BATCH_SIZE = 1 20 21 # --- 1. データセットの準備 --- 22 # JSON Linesファイルをロード 23 dataset = load_dataset("json", data_files=DATASET_PATH, split="train") 24 25 26 def format_prompt(example): 27 """ 28 Typstコードを入力（Instruction）として、JSON出力を期待するプロンプト形式に整形する。 29 MistralのInstructionフォーマット（[INST]...[/INST]）を使用。 30 """ 31 32 # 指示: Typstコードをレビューし、JSON形式で指摘事項を出力せよ 33 prompt = f"""[INST] 34 以下のTypstコードのレビューを行い、文書のフォーマット、論理、文法に関するすべての指摘事項を、定義されたJSON形式のリストとして出力してください。 35 36 Typstコード: 37 --- 38 {example["instruction"]} 39 --- 40 41 JSON形式の指摘リスト: 42 [/INST]""" 43 44 # LLMの学習データは 'Instruction' と 'Output' を結合した形式にする 45 example["text"] = prompt + example["output"] 46 return example 47 48 49 # データセットをプロンプト形式に変換 50 dataset = dataset.map(format_prompt, remove_columns=["instruction", "output"]) 51 52 # トークナイザをロード（修正済みセクション） 53 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) 54 tokenizer.pad_token = tokenizer.eos_token 55 56 57 def tokenize_function(examples): 58 tokenized = tokenizer( 59 examples["text"], truncation=True, padding="max_length", max_length=1024 60 ) 61 # 【✅ 修正箇所】 62 # input_idsを複製してlabelsとする 63 # Data Collatorがinput_ids（モデル入力）とlabels（教師データ）の両方を見つけられるようにする 64 tokenized["labels"] = tokenized["input_ids"].copy() 65 return tokenized 66 67 68 tokenized_dataset = dataset.map( 69 tokenize_function, batched=True, remove_columns=["text"] 70 ) 71 72 # --- 2. 量子化設定とモデルのロード --- 73 # QLoRAのための4bit量子化設定 74 bnb_config = BitsAndBytesConfig( 75 load_in_4bit=True, 76 bnb_4bit_quant_type="nf4", 77 bnb_4bit_compute_dtype=torch.bfloat16, 78 bnb_4bit_use_double_quant=False, 79 ) 80 81 # モデルとトークナイザをロード 82 model = AutoModelForCausalLM.from_pretrained( 83 MODEL_NAME, 84 quantization_config=bnb_config, 85 device_map="auto", 86 ) 87 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) 88 tokenizer.pad_token = tokenizer.eos_token # トークナイザ設定 89 90 # 4bit学習のためのモデル準備 91 model.gradient_checkpointing_enable() 92 model = prepare_model_for_kbit_training(model) 93 94 95 # --- 3. LoRA設定とモデル適用 --- 96 lora_config = LoraConfig( 97 r=LORA_R, 98 lora_alpha=LORA_ALPHA, 99 lora_dropout=LORA_DROPOUT, 100 bias="none", 101 task_type="CAUSAL_LM", # テキスト生成タスク 102 target_modules=[ 103 "q_proj", 104 "k_proj", 105 "v_proj", 106 "o_proj", 107 ], # Mistralの主要なAttention層 108 ) 109 110 # PEFTモデルの取得 111 model = get_peft_model(model, lora_config) 112 model.print_trainable_parameters() # 学習対象のパラメータ数を確認（非常に少ないはず） 113 114 115 # --- 4. 学習の実行 --- 116 # 【✅ 修正 2: TrainingArgumentsの設定変更】 117 training_args = TrainingArguments( 118 output_dir=OUTPUT_DIR, 119 per_device_train_batch_size=BATCH_SIZE, 120 gradient_accumulation_steps=4, 121 warmup_steps=10, 122 max_steps=20, 123 learning_rate=2e-4, 124 fp16=True, 125 logging_steps=5, 126 # 【ここを修正】データセットの列名がTrainerと完全に一致しなくても許容する 127 remove_unused_columns=False, 128 ) 129 130 trainer = Trainer( 131 model=model, 132 # 【修正 3: トークン化後のデータセットを使用】 133 train_dataset=tokenized_dataset, 134 args=training_args, 135 tokenizer=tokenizer, 136 ) 137 138 # 学習開始 139 trainer.train() 140 141 # 学習済みアダプターの保存 142 trainer.model.save_pretrained(OUTPUT_DIR)