ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter8/continued-pretraining/continued-pretrain.py

@@ -0,0 +1,645 @@
+# -*- coding: utf-8 -*-
+"""Continued pretraining - Korean + Unsloth.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1tEd1FrOXWMnCU9UIvdYhs61tkxdMuKZu
+
+To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4 Google Colab instance!
+<div class="align-center">
+  <a href="https://github.com/unslothai/unsloth"><img src="https://github.com/unslothai/unsloth/raw/main/images/unsloth%20new%20logo.png" width="115"></a>
+  <a href="https://discord.gg/u54VK8m8tk"><img src="https://github.com/unslothai/unsloth/raw/main/images/Discord button.png" width="145"></a>
+  <a href="https://ko-fi.com/unsloth"><img src="https://github.com/unslothai/unsloth/raw/main/images/Kofi button.png" width="145"></a></a> Join Discord if you need help + ⭐ <i>Star us on <a href="https://github.com/unslothai/unsloth">Github</a> </i> ⭐
+</div>
+
+To install Unsloth on your own computer, follow the installation instructions on our Github page [here](https://github.com/unslothai/unsloth#installation-instructions---conda).
+
+You will learn how to do [data prep](#Data), how to [train](#Train), how to [run the model](#Inference), & [how to save it](#Save) (eg for Llama.cpp).
+
+We will use the Korean subset of the [Wikipedia dataset](https://huggingface.co/datasets/wikimedia/wikipedia) to first continually pretrain Mistral v3, then use the [Alpaca GPT4 Dataset](https://huggingface.co/datasets/FreedomIntelligence/alpaca-gpt4-korean) translated into Korean to further finetune the model to let it follow instructions in Korean.
+"""
+
+# ============================================================================
+# IMPORTS
+# ============================================================================
+import os
+import argparse
+
+
+# ============================================================================
+# 命令行参数（默认值与原始脚本硬编码值完全一致，保证行为不变）
+# ============================================================================
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="韩语 Mistral 继续预训练 + 指令微调（Unsloth / LoRA）。"
+                    "先用韩语维基百科做继续预训练注入韩语能力，再用韩语 Alpaca 数据做 SFT。",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    # 基础模型
+    parser.add_argument("--base_model", type=str, default="unsloth/mistral-7b-v0.3",
+                        help="基础模型名称（HuggingFace / Unsloth 仓库名）")
+    parser.add_argument("--max_seq_len", type=int, default=2048,
+                        help="最大序列长度")
+    parser.add_argument("--no_4bit", action="store_true",
+                        help="关闭 4bit 量化加载（默认开启 4bit 以节省显存）")
+
+    # LoRA 配置
+    parser.add_argument("--lora_rank", type=int, default=128,
+                        help="LoRA 秩 r（继续预训练建议较大，如 128）")
+    parser.add_argument("--lora_alpha", type=int, default=32,
+                        help="LoRA alpha")
+
+    # 数据集
+    parser.add_argument("--wiki_dataset", type=str, default="wikimedia/wikipedia",
+                        help="继续预训练用的维基百科数据集")
+    parser.add_argument("--wiki_config", type=str, default="20231101.ko",
+                        help="维基百科数据集的语言/版本配置（默认韩语快照）")
+    parser.add_argument("--wiki_train_size", type=float, default=0.05,
+                        help="维基百科数据集抽样比例（0~1，默认取 5%% 以加速训练）")
+    parser.add_argument("--alpaca_dataset", type=str,
+                        default="FreedomIntelligence/alpaca-gpt4-korean",
+                        help="指令微调（SFT）用的韩语 Alpaca 数据集")
+
+    # 训练超参数
+    parser.add_argument("--pretrain_epochs", type=int, default=1,
+                        help="继续预训练阶段的训练轮数")
+    parser.add_argument("--pretrain_max_steps", type=int, default=-1,
+                        help="继续预训练最大步数（-1 表示不限制，按 epoch 训练）")
+    parser.add_argument("--sft_epochs", type=int, default=2,
+                        help="指令微调阶段的训练轮数")
+    parser.add_argument("--sft_max_steps", type=int, default=-1,
+                        help="指令微调最大步数（-1 表示不限制，按 epoch 训练）")
+
+    # 输出目录
+    parser.add_argument("--pretrain_output_dir", type=str, default="outputs_pretrain",
+                        help="继续预训练的检查点目录")
+    parser.add_argument("--sft_output_dir", type=str, default="outputs_sft",
+                        help="指令微调的检查点目录")
+    parser.add_argument("--pretrained_save_dir", type=str, default="lora_model_pretrained",
+                        help="继续预训练后保存的 LoRA 模型目录")
+    parser.add_argument("--final_save_dir", type=str, default="lora_model",
+                        help="指令微调后保存的最终 LoRA 模型目录")
+
+    return parser.parse_args()
+
+
+# 先解析参数（放在重型导入之前，这样 `--help` 无需 GPU / Unsloth 也能运行）
+args = parse_args()
+
+import torch
+from unsloth import FastLanguageModel, is_bfloat16_supported, UnslothTrainer, UnslothTrainingArguments
+from transformers import TrainingArguments, TextStreamer
+from datasets import load_dataset
+
+# ANSI color codes for colored output
+class Colors:
+    HEADER = '\033[95m'
+    BLUE = '\033[94m'
+    CYAN = '\033[96m'
+    GREEN = '\033[92m'
+    YELLOW = '\033[93m'
+    RED = '\033[91m'
+    ENDC = '\033[0m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+
+def print_section(title, color=Colors.CYAN):
+    """Print a colored section header"""
+    print(f"\n{color}{Colors.BOLD}{'='*70}")
+    print(f"{title}")
+    print(f"{'='*70}{Colors.ENDC}\n")
+
+# ============================================================================
+# MODEL SETUP
+# ============================================================================
+print_section("🚀 LOADING MODEL", Colors.BLUE)
+
+max_seq_length = args.max_seq_len # Choose any! We auto support RoPE Scaling internally!
+dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+load_in_4bit = not args.no_4bit # Use 4bit quantization to reduce memory usage. Can be False.
+
+# 4bit pre quantized models we support for 4x faster downloading + no OOMs.
+fourbit_models = [
+    "unsloth/mistral-7b-v0.3-bnb-4bit",      # New Mistral v3 2x faster!
+    "unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
+    "unsloth/llama-3-8b-bnb-4bit",           # Llama-3 15 trillion tokens model 2x faster!
+    "unsloth/llama-3-8b-Instruct-bnb-4bit",
+    "unsloth/llama-3-70b-bnb-4bit",
+    "unsloth/Phi-3-mini-4k-instruct",        # Phi-3 2x faster!
+    "unsloth/Phi-3-medium-4k-instruct",
+    "unsloth/mistral-7b-bnb-4bit",
+    "unsloth/gemma-7b-bnb-4bit",             # Gemma 2.2x faster!
+] # More models at https://huggingface.co/unsloth
+
+print(f"{Colors.GREEN}Loading {args.base_model}...{Colors.ENDC}")
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name = args.base_model, # Choose ANY! eg teknium/OpenHermes-2.5-Mistral-7B
+    max_seq_length = max_seq_length,
+    dtype = dtype,
+    load_in_4bit = load_in_4bit,
+    # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
+)
+
+print_section("⚙️ ADDING LORA ADAPTERS", Colors.BLUE)
+print(f"{Colors.GREEN}Adding LoRA adapters - only updating 1-10% of parameters!")
+print(f"Including embed_tokens and lm_head for continual pretraining{Colors.ENDC}")
+
+model = FastLanguageModel.get_peft_model(
+    model,
+    r = args.lora_rank, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                      "gate_proj", "up_proj", "down_proj",
+                      "embed_tokens", "lm_head",], # Add for continual pretraining
+    lora_alpha = args.lora_alpha,
+    lora_dropout = 0, # Supports any, but = 0 is optimized
+    bias = "none",    # Supports any, but = "none" is optimized
+    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
+    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
+    random_state = 3407,
+    use_rslora = True,   # We support rank stabilized LoRA
+    loftq_config = None, # And LoftQ
+)
+
+# ============================================================================
+# TESTING ORIGINAL MODEL (BASELINE)
+# ============================================================================
+print_section("🧪 TESTING ORIGINAL MODEL (BASELINE)", Colors.CYAN)
+print(f"{Colors.YELLOW}Testing the base Mistral model BEFORE any training")
+print(f"This establishes baseline for Korean and English capabilities{Colors.ENDC}")
+
+FastLanguageModel.for_inference(model)
+text_streamer = TextStreamer(tokenizer)
+
+# Prepare prompts (define them early)
+_wikipedia_prompt = """Wikipedia Article
+### Title: {}
+
+### Article:
+{}"""
+
+wikipedia_prompt_korean = """위키피디아 기사
+### 제목: {}
+
+### 기사:
+{}"""
+
+_alpaca_prompt_english = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
+
+### Instruction:
+{}
+
+### Response:
+{}"""
+
+alpaca_prompt_korean = """다음은 작업을 설명하는 명령입니다. 요청을 적절하게 완료하는 응답을 작성하세요.
+
+### 지침:
+{}
+
+### 응답:
+{}"""
+
+# Test 1: Korean Wikipedia (baseline - should be poor)
+print(f"\n{Colors.BOLD}Test 1: Korean Wikipedia Article (인공지능){Colors.ENDC}")
+print("="*70)
+test_prompt = wikipedia_prompt_korean.format("인공지능", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 2: English Wikipedia (baseline - should be good)
+print(f"\n{Colors.BOLD}Test 2: English Wikipedia Article (Artificial Intelligence){Colors.ENDC}")
+print("="*70)
+test_prompt = _wikipedia_prompt.format("Artificial Intelligence", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 3: Korean Instruction (baseline - should be poor)
+print(f"\n{Colors.BOLD}Test 3: Korean Instruction (Korean Culture){Colors.ENDC}")
+print("="*70)
+test_prompt = alpaca_prompt_korean.format("한국의 전통 음식인 김치에 대해 설명하세요.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 4: English Instruction (baseline - should be good)
+print(f"\n{Colors.BOLD}Test 4: English Instruction (American Culture){Colors.ENDC}")
+print("="*70)
+test_prompt = _alpaca_prompt_english.format("Explain about Thanksgiving turkey, a traditional American food.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 5: Korean Instruction - Seoul (baseline - should be poor)
+print(f"\n{Colors.BOLD}Test 5: Korean Instruction (Korean Geography){Colors.ENDC}")
+print("="*70)
+test_prompt = alpaca_prompt_korean.format("대한민국의 수도인 서울에 대해 간단히 소개해주세요.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+print(f"{Colors.GREEN}✓ Baseline testing complete. Original model should be good at English but poor at Korean.{Colors.ENDC}\n")
+
+# ============================================================================
+# DATA PREPARATION
+# ============================================================================
+print_section("📚 DATA PREPARATION - WIKIPEDIA KOREAN DATASET", Colors.CYAN)
+print(f"{Colors.YELLOW}Loading Korean Wikipedia dataset...")
+print(f"We'll use 5% of the dataset to speed up training{Colors.ENDC}")
+
+# Use the prompts already defined above
+EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
+def formatting_prompts_func_wiki(examples):
+    titles = examples["title"]
+    texts  = examples["text"]
+    outputs = []
+    for title, text in zip(titles, texts):
+        # Must add EOS_TOKEN, otherwise your generation will go on forever!
+        text = wikipedia_prompt_korean.format(title, text) + EOS_TOKEN
+        outputs.append(text)
+    return { "text" : outputs, }
+
+dataset = load_dataset(args.wiki_dataset, args.wiki_config, split = "train",)
+# We select 5% of the data to make training faster!
+dataset = dataset.train_test_split(train_size = args.wiki_train_size)["train"]
+dataset = dataset.map(formatting_prompts_func_wiki, batched = True,)
+
+print(f"{Colors.GREEN}✓ Wikipedia dataset loaded: {len(dataset)} examples{Colors.ENDC}")
+
+print_section("📚 DATA PREPARATION - ALPACA KOREAN DATASET", Colors.CYAN)
+print(f"{Colors.YELLOW}Loading Alpaca GPT4 Korean dataset for instruction finetuning...{Colors.ENDC}")
+
+alpaca_dataset = load_dataset(args.alpaca_dataset, split = "train")
+
+# Use the prompts already defined above
+def formatting_prompts_func_alpaca(conversations):
+    texts = []
+    conversations = conversations["conversations"]
+    for convo in conversations:
+        # Must add EOS_TOKEN, otherwise your generation will go on forever!
+        text = alpaca_prompt_korean.format(convo[0]["value"], convo[1]["value"]) + EOS_TOKEN
+        texts.append(text)
+    return { "text" : texts, }
+
+alpaca_dataset = alpaca_dataset.map(formatting_prompts_func_alpaca, batched = True,)
+
+print(f"{Colors.GREEN}✓ Alpaca dataset loaded: {len(alpaca_dataset)} examples{Colors.ENDC}")
+print(f"\nExample from Alpaca dataset:")
+print(alpaca_dataset[0])
+
+# ============================================================================
+# CONTINUED PRETRAINING
+# ============================================================================
+print_section("🎯 CONTINUED PRETRAINING ON KOREAN WIKIPEDIA", Colors.GREEN)
+print(f"{Colors.YELLOW}Training the model on Korean Wikipedia to learn the language...")
+print(f"Using embedding_learning_rate (1e-5) smaller than learning_rate (5e-5)")
+print(f"💾 Checkpoints will be saved every 100 steps to: outputs_pretrain/")
+print(f"Only the 5 most recent checkpoints will be kept{Colors.ENDC}")
+
+# Set WandB project for pretraining
+os.environ["WANDB_PROJECT"] = "unsloth-continued-pretraining"
+
+trainer = UnslothTrainer(
+    model = model,
+    tokenizer = tokenizer,
+    train_dataset = dataset,
+    dataset_text_field = "text",
+    max_seq_length = max_seq_length,
+    dataset_num_proc = 2,
+
+    args = UnslothTrainingArguments(
+        per_device_train_batch_size = 2,
+        gradient_accumulation_steps = 8,
+
+        # Use warmup_ratio and num_train_epochs for longer runs!
+        max_steps = args.pretrain_max_steps,
+        warmup_steps = 10,
+        warmup_ratio = 0.1,
+        num_train_epochs = args.pretrain_epochs,
+
+        # Select a 2 to 10x smaller learning rate for the embedding matrices!
+        learning_rate = 5e-5,
+        embedding_learning_rate = 1e-5,
+
+        fp16 = not is_bfloat16_supported(),
+        bf16 = is_bfloat16_supported(),
+        logging_steps = 1,
+        optim = "adamw_8bit",
+        weight_decay = 0.01,
+        lr_scheduler_type = "linear",
+        seed = 42,
+        output_dir = args.pretrain_output_dir,
+        
+        # Checkpoint saving
+        save_strategy = "steps",
+        save_steps = 100,
+        # Keep only the 5 most recent checkpoints, as the banner above
+        # promises. Without this HF keeps every one, and a ~2000-step run
+        # writes ~20 LoRA+optimizer checkpoints (fills a free-Colab disk).
+        save_total_limit = 5,
+        
+        # WandB configuration
+        report_to = "wandb",
+        run_name = "korean-mistral-pretrain",
+    ),
+)
+
+print_section("💾 MEMORY STATS - BEFORE PRETRAINING", Colors.YELLOW)
+gpu_stats = torch.cuda.get_device_properties(0)
+start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
+max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
+print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
+print(f"{start_gpu_memory} GB of memory reserved.")
+
+trainer_stats = trainer.train()
+
+# Finish wandb run for pretraining
+import wandb
+if wandb.run is not None:
+    wandb.finish()
+    print(f"{Colors.CYAN}✓ Finished wandb run for pretraining{Colors.ENDC}")
+
+print_section("💾 SAVING PRETRAINED MODEL", Colors.GREEN)
+model.save_pretrained(args.pretrained_save_dir) # Local saving
+tokenizer.save_pretrained(args.pretrained_save_dir)
+print(f"{Colors.GREEN}✓ Model saved to: {args.pretrained_save_dir}/{Colors.ENDC}")
+
+# ============================================================================
+# TESTING PRETRAINED MODEL
+# ============================================================================
+print_section("🧪 TESTING PRETRAINED MODEL (AFTER KOREAN PRETRAINING)", Colors.CYAN)
+print(f"{Colors.YELLOW}Testing after Korean pretraining - before instruction finetuning")
+print(f"Korean should improve, English should remain strong{Colors.ENDC}")
+
+# Test the pretrained model before instruction finetuning
+FastLanguageModel.for_inference(model) # Enable native 2x faster inference
+text_streamer = TextStreamer(tokenizer)
+
+# Test 1: Korean Wikipedia (same as baseline)
+print(f"\n{Colors.BOLD}Test 1: Korean Wikipedia Article (인공지능){Colors.ENDC}")
+print("="*70)
+test_prompt = wikipedia_prompt_korean.format("인공지능", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 2: English Wikipedia (same as baseline)
+print(f"\n{Colors.BOLD}Test 2: English Wikipedia Article (Artificial Intelligence){Colors.ENDC}")
+print("="*70)
+test_prompt = _wikipedia_prompt.format("Artificial Intelligence", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 3: Korean Instruction (same as baseline - should improve but not follow perfectly yet)
+print(f"\n{Colors.BOLD}Test 3: Korean Instruction (Korean Culture){Colors.ENDC}")
+print("="*70)
+test_prompt = alpaca_prompt_korean.format("한국의 전통 음식인 김치에 대해 설명하세요.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 4: English Instruction (same as baseline)
+print(f"\n{Colors.BOLD}Test 4: English Instruction (American Culture){Colors.ENDC}")
+print("="*70)
+test_prompt = _alpaca_prompt_english.format("Explain about Thanksgiving turkey, a traditional American food.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 5: Korean Instruction - Seoul (should improve but not perfect yet)
+print(f"\n{Colors.BOLD}Test 5: Korean Instruction (Korean Geography){Colors.ENDC}")
+print("="*70)
+test_prompt = alpaca_prompt_korean.format("대한민국의 수도인 서울에 대해 간단히 소개해주세요.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+print(f"{Colors.GREEN}✓ Pretrained model testing complete.")
+print(f"{Colors.CYAN}💡 Korean should be significantly improved, English should remain strong.")
+print(f"Instruction following may improve but not perfect yet - that's what SFT is for.{Colors.ENDC}\n")
+
+# ============================================================================
+# INSTRUCTION FINETUNING
+# ============================================================================
+print_section("🎓 INSTRUCTION FINETUNING ON ALPACA KOREAN", Colors.GREEN)
+print(f"{Colors.YELLOW}Now finetuning the model to follow Korean instructions...")
+print(f"Using the Alpaca GPT4 dataset translated to Korean")
+print(f"💾 Checkpoints will be saved every 100 steps to: outputs_sft/")
+print(f"Only the 5 most recent checkpoints will be kept{Colors.ENDC}")
+
+# Set WandB project for finetuning
+os.environ["WANDB_PROJECT"] = "unsloth-continued-finetuning"
+
+trainer = UnslothTrainer(
+    model = model,
+    tokenizer = tokenizer,
+    train_dataset = alpaca_dataset,
+    dataset_text_field = "text",
+    max_seq_length = max_seq_length,
+    dataset_num_proc = 8,
+
+    args = UnslothTrainingArguments(
+        per_device_train_batch_size = 2,
+        gradient_accumulation_steps = 8,
+
+        # Use num_train_epochs and warmup_ratio for longer runs!
+        max_steps = args.sft_max_steps,
+        warmup_steps = 10,
+        warmup_ratio = 0.1,
+        num_train_epochs = args.sft_epochs,
+
+        # Select a 2 to 10x smaller learning rate for the embedding matrices!
+        learning_rate = 5e-5,
+        embedding_learning_rate = 1e-5,
+
+        fp16 = not is_bfloat16_supported(),
+        bf16 = is_bfloat16_supported(),
+        logging_steps = 1,
+        optim = "adamw_8bit",
+        weight_decay = 0.00,
+        lr_scheduler_type = "linear",
+        seed = 42,
+        output_dir = args.sft_output_dir,
+        
+        # Checkpoint saving
+        save_strategy = "steps",
+        save_steps = 100,
+        # Keep only the 5 most recent checkpoints, as the banner above
+        # promises. Without this HF keeps every one, and a ~2000-step run
+        # writes ~20 LoRA+optimizer checkpoints (fills a free-Colab disk).
+        save_total_limit = 5,
+        
+        # WandB configuration
+        report_to = "wandb",
+        run_name = "korean-mistral-finetune",
+    ),
+)
+
+trainer_stats = trainer.train()
+
+# Finish wandb run for SFT
+if wandb.run is not None:
+    wandb.finish()
+    print(f"{Colors.CYAN}✓ Finished wandb run for SFT{Colors.ENDC}")
+
+print_section("📊 FINAL MEMORY AND TIME STATS", Colors.YELLOW)
+used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
+used_memory_for_lora = round(used_memory - start_gpu_memory, 3)
+used_percentage = round(used_memory         /max_memory*100, 3)
+lora_percentage = round(used_memory_for_lora/max_memory*100, 3)
+print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
+print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.")
+print(f"Peak reserved memory = {used_memory} GB.")
+print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
+print(f"Peak reserved memory % of max memory = {used_percentage} %.")
+print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
+
+# ============================================================================
+# INFERENCE - TESTING FINETUNED MODEL
+# ============================================================================
+print_section("🎯 INFERENCE - TESTING FINETUNED MODEL", Colors.CYAN)
+print(f"{Colors.YELLOW}Testing the instruction-finetuned model - should follow instructions in Korean AND English{Colors.ENDC}")
+
+FastLanguageModel.for_inference(model) # Enable native 2x faster inference
+text_streamer = TextStreamer(tokenizer)
+
+# Test 1: Korean Wikipedia (same as baseline)
+print(f"\n{Colors.BOLD}Test 1: Korean Wikipedia Article (인공지능){Colors.ENDC}")
+print("="*70)
+test_prompt = wikipedia_prompt_korean.format("인공지능", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 2: English Wikipedia (same as baseline)
+print(f"\n{Colors.BOLD}Test 2: English Wikipedia Article (Artificial Intelligence){Colors.ENDC}")
+print("="*70)
+test_prompt = _wikipedia_prompt.format("Artificial Intelligence", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 3: Korean Instruction (same as baseline - should follow WELL now)
+print(f"\n{Colors.BOLD}Test 3: Korean Instruction (Korean Culture){Colors.ENDC}")
+print("="*70)
+test_prompt = alpaca_prompt_korean.format("한국의 전통 음식인 김치에 대해 설명하세요.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 4: English Instruction (same as baseline)
+print(f"\n{Colors.BOLD}Test 4: English Instruction (American Culture){Colors.ENDC}")
+print("="*70)
+test_prompt = _alpaca_prompt_english.format("Explain about Thanksgiving turkey, a traditional American food.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+# Test 5: Korean Instruction - Seoul (should follow WELL now)
+print(f"\n{Colors.BOLD}Test 5: Korean Instruction (Korean Geography){Colors.ENDC}")
+print("="*70)
+test_prompt = alpaca_prompt_korean.format("대한민국의 수도인 서울에 대해 간단히 소개해주세요.", "")
+inputs = tokenizer([test_prompt], return_tensors = "pt").to("cuda")
+_ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 150, use_cache = True)
+print("="*70 + "\n")
+
+print(f"{Colors.GREEN}✓ Finetuned model testing complete!")
+print(f"{Colors.CYAN}💡 Model should now follow instructions well in both Korean AND English.{Colors.ENDC}\n")
+
+# ============================================================================
+# SAVING FINAL MODEL
+# ============================================================================
+print_section("💾 SAVING FINAL FINETUNED MODEL", Colors.GREEN)
+
+model.save_pretrained(args.final_save_dir) # Local saving
+tokenizer.save_pretrained(args.final_save_dir)
+print(f"{Colors.GREEN}✓ Final model saved to: {args.final_save_dir}/{Colors.ENDC}")
+# model.push_to_hub("your_name/lora_model", token = "...") # Online saving
+# tokenizer.push_to_hub("your_name/lora_model", token = "...") # Online saving
+
+print(f"\n{Colors.CYAN}Note: This only saves LoRA adapters, not the full model.")
+print(f"For 16bit or GGUF formats, see the export options below.{Colors.ENDC}")
+
+# ============================================================================
+# LOADING SAVED MODEL (OPTIONAL)
+# ============================================================================
+"""
+print_section("📥 LOADING SAVED MODEL", Colors.BLUE)
+
+# Set to True to test loading
+if False:
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = "lora_model", # YOUR MODEL YOU USED FOR TRAINING
+        max_seq_length = max_seq_length,
+        dtype = dtype,
+        load_in_4bit = load_in_4bit,
+    )
+    FastLanguageModel.for_inference(model) # Enable native 2x faster inference
+
+    inputs = tokenizer(
+    [
+        alpaca_prompt.format(
+            # "Describe the planet Earth extensively.", # instruction
+            "지구를 광범위하게 설명하세요.",
+            "", # output - leave this blank for generation!
+        ),
+    ], return_tensors = "pt").to("cuda")
+
+    text_streamer = TextStreamer(tokenizer)
+    _ = model.generate(**inputs, streamer = text_streamer, max_new_tokens = 256)
+"""
+
+# ============================================================================
+# EXPORT OPTIONS
+# ============================================================================
+print_section("📦 EXPORT OPTIONS", Colors.BLUE)
+print(f"{Colors.YELLOW}Various export formats available (currently disabled):{Colors.ENDC}")
+print("• Float16 merged model")
+print("• 4bit merged model")
+print("• LoRA adapters only")
+print("• GGUF format for llama.cpp")
+
+# Merge to 16bit
+if False: 
+    print_section("💾 EXPORTING TO FLOAT16", Colors.GREEN)
+    model.save_pretrained_merged("model", tokenizer, save_method = "merged_16bit",)
+    # model.push_to_hub_merged("hf/model", tokenizer, save_method = "merged_16bit", token = "")
+
+# Merge to 4bit
+if False: 
+    print_section("💾 EXPORTING TO 4BIT", Colors.GREEN)
+    model.save_pretrained_merged("model", tokenizer, save_method = "merged_4bit",)
+    # model.push_to_hub_merged("hf/model", tokenizer, save_method = "merged_4bit", token = "")
+
+# Just LoRA adapters
+if False: 
+    print_section("💾 EXPORTING LORA ADAPTERS", Colors.GREEN)
+    model.save_pretrained_merged("model", tokenizer, save_method = "lora",)
+    # model.push_to_hub_merged("hf/model", tokenizer, save_method = "lora", token = "")
+
+# GGUF exports
+if False:
+    print_section("💾 EXPORTING TO GGUF Q8_0", Colors.GREEN)
+    model.save_pretrained_gguf("model", tokenizer,)
+    # model.push_to_hub_gguf("hf/model", tokenizer, token = "")
+
+if False:
+    print_section("💾 EXPORTING TO GGUF F16", Colors.GREEN)
+    model.save_pretrained_gguf("model", tokenizer, quantization_method = "f16")
+    # model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "f16", token = "")
+
+if False:
+    print_section("💾 EXPORTING TO GGUF Q4_K_M", Colors.GREEN)
+    model.save_pretrained_gguf("model", tokenizer, quantization_method = "q4_k_m")
+    # model.push_to_hub_gguf("hf/model", tokenizer, quantization_method = "q4_k_m", token = "")
+
+print_section("✅ TRAINING COMPLETE!", Colors.GREEN)
+print(f"{Colors.BOLD}Your Korean Mistral model is ready to use!{Colors.ENDC}")
+print(f"\n{Colors.CYAN}For more information:{Colors.ENDC}")
+print("• Discord: https://discord.gg/u54VK8m8tk")
+print("• GitHub: https://github.com/unslothai/unsloth")
+print("• Documentation: https://docs.unsloth.ai")