continued-pretraining

第7章 · 模型后训练 · 配套项目 chapter7/continued-pretraining

项目说明

继续预训练：让模型学会一门新语言（韩语 Mistral）

本目录对应《深入理解 AI Agent》第 7 章 实验 7-5 ★★：继续预训练学习新语言。

项目简介

以 Mistral 7B v0.3 为基础模型（主要用英语预训练，对韩语几乎没有理解能力），通过韩语维基百科继续预训练注入韩语能力，再用韩语指令数据做 SFT，最终得到一个既能理解韩语、又能用韩语遵循指令的模型。

本实验想说明的核心观点：要让模型记住大量新领域知识（这里是一门新语言），靠的是继续预训练，而不是 SFT。 模型在预训练阶段已经具备通用的语言建模能力，继续预训练只是让它适应新的数据分布，成本远低于从头训练。

整个流程分两个阶段：

1. 继续预训练（Continued Pretraining）：在韩语维基百科上做无监督的“预测下一个词”训练，让模型学会韩语的词汇与句法。
2. 指令微调（SFT）：在韩语 Alpaca 指令数据上训练，让模型学会“用韩语遵循指令”。

一个关键工程点是缓解灾难性遗忘（Catastrophic Forgetting）：学了新语言不能把原来的英语能力忘掉。书中讨论的通用做法是用混合数据（约 80% 目标语言 + 20% 原语言）来平衡；本实现则采用 LoRA + 训练 embed_tokens/lm_head 的参数高效方案——只更新适配器与词嵌入，基础权重保持不变，从而在注入韩语的同时尽量保留英语。评测结果（见下文）显示英语能力基本得到保留。

目录结构

continued-pretraining/
├── README.md                 # 本文档
├── continued-pretrain.py     # 训练主脚本：继续预训练 + SFT，产出两个 LoRA 模型
├── evaluate_model.py         # 单模型评测：在韩英任务上生成样例
├── compare_models.py         # 三阶段对比：基础 → 继续预训练 → 指令微调 并排生成
├── model_eval_results.md     # 真实运行的完整评测输出与结论（RTX 4090）
├── README_EVALUATION.md      # 评测脚本的详细用法说明
└── requirements.txt          # 依赖清单

训练脚本运行后会产出两个本地目录（仅保存 LoRA 适配器，不含完整模型）：

• lora_model_pretrained/：继续预训练之后、SFT 之前的模型
• lora_model/：最终指令微调之后的模型

系统要求与依赖

• GPU：需要支持 CUDA 的 NVIDIA GPU。默认以 4bit 量化加载 Mistral-7B，可在约 24GB 显存的消费级显卡（如 RTX 4090）上完成训练，model_eval_results.md 中的结果即在 RTX 4090 上产出。
• 框架：Unsloth（高效 LoRA 训练）、PyTorch、Transformers、Datasets、bitsandbytes。
• 可选：wandb（实验跟踪，脚本默认 report_to="wandb"）。

pip install -r requirements.txt

注意：Unsloth 依赖 GPU 与匹配的 CUDA/PyTorch 版本，无法在纯 CPU 环境下训练或推理。各脚本的 --help 已做延迟导入，可在没有 GPU 的机器上直接查看参数说明。

快速开始

1. 训练（继续预训练 + SFT）

用默认超参数一键完成两个阶段（韩语维基百科 5% 子集做继续预训练，随后用韩语 Alpaca 做 SFT）：

python continued-pretrain.py

脚本会依次：加载基础模型 → 打印基线测试 → 韩语维基继续预训练 → 保存 lora_model_pretrained/ → 韩语指令 SFT → 保存 lora_model/。

常用参数（默认值与脚本原始硬编码一致，改动才会偏离原实验）：

python continued-pretrain.py \
    --base_model unsloth/mistral-7b-v0.3 \
    --wiki_config 20231101.ko \
    --wiki_train_size 0.05 \
    --alpaca_dataset FreedomIntelligence/alpaca-gpt4-korean \
    --lora_rank 128 \
    --max_seq_len 2048 \
    --pretrain_epochs 1 \
    --sft_epochs 2 \
    --pretrained_save_dir lora_model_pretrained \
    --final_save_dir lora_model

• 想快速冒烟测试，可用 --pretrain_max_steps 20 --sft_max_steps 20 只跑很少的步数。
• 想换一门语言：把 --wiki_config 换成对应维基快照（如 20231101.ja 日语）、--alpaca_dataset 换成对应语言的指令集即可。
• 完整参数见 python continued-pretrain.py --help。

2. 评测单个模型

# 评测最终微调模型（默认加载 lora_model/）
python evaluate_model.py

# 评测继续预训练后、SFT 前的模型
python evaluate_model.py --pretrained

# 生成更长、使用采样
python evaluate_model.py --max_new_tokens 300 --use_sampling --temperature 0.7

更多用法详见 README_EVALUATION.md。

3. 三阶段并排对比

同时加载基础模型 / 继续预训练模型 / 指令微调模型，在同一组中韩英提示上并排生成，直观展示韩语能力的提升与英语能力的保留：

python compare_models.py

# 指定模型目录与生成参数
python compare_models.py \
    --pretrained_path lora_model_pretrained \
    --finetuned_path lora_model \
    --max_new_tokens 150 \
    --temperature 0.3

实验结果

真实运行的完整输出、逐条对比与结论见 model_eval_results.md（在 RTX 4090 上产出，请以该文件中的实际结果为准，本文不再复述具体数值）。其主要结论可概括为：

• 方法论成立：继续预训练 + SFT 确实能为模型注入新语言能力，韩语从“基本不可用”提升到“流畅、能遵循指令”，同时英语能力基本得到保留（无明显灾难性遗忘）。
• 数据质量是瓶颈：仅用 5% 的维基百科语料，通用语言能力提升明显，但特定文化知识（如泡菜等）仍常出错——说明对具体知识域而言，数据覆盖与质量比训练方法本身更关键。

参考资料

• Unsloth 文档：https://docs.unsloth.ai
• 基础模型：unsloth/mistral-7b-v0.3
• 继续预训练语料：wikimedia/wikipedia（20231101.ko）
• 指令微调语料：FreedomIntelligence/alpaca-gpt4-korean

源代码

compare_models.py

# -*- coding: utf-8 -*-
"""
Compare baseline → pretrained → finetuned Korean Mistral models (3-way comparison)
Shows progression from original model to final Korean-capable model
"""

import argparse

# 说明：unsloth / torch 等重型依赖在函数内按需导入，
# 这样 `python compare_models.py --help` 无需 GPU 环境即可查看参数。

# 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}{'='*80}")
    print(f"{title}")
    print(f"{'='*80}{Colors.ENDC}\n")

def load_baseline_model(base_model="unsloth/mistral-7b-v0.3", max_seq_length=2048):
    """Load the original Mistral model (before any training)"""
    from unsloth import FastLanguageModel
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=base_model,
        max_seq_length=max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )
    FastLanguageModel.for_inference(model)
    return model, tokenizer

def load_model(model_path, max_seq_length=2048):
    """Load a trained LoRA model"""
    from unsloth import FastLanguageModel
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_path,
        max_seq_length=max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )
    FastLanguageModel.for_inference(model)
    return model, tokenizer

def generate_text(model, tokenizer, prompt, max_new_tokens=150, temperature=0.3):
    """Generate text without streaming"""
    inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        use_cache=True,
        do_sample=True,
        temperature=temperature,
        pad_token_id=tokenizer.eos_token_id,
    )
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Remove the prompt from the output
    response = generated_text[len(prompt):].strip()
    return response

def compare_on_prompt(baseline_model, baseline_tokenizer,
                     pretrained_model, pretrained_tokenizer, 
                     finetuned_model, finetuned_tokenizer,
                     prompt, test_name, prompt_translation=None,
                     max_new_tokens=150, temperature=0.3):
    """Compare three models on the same prompt"""

    print(f"\n{Colors.BOLD}{'='*80}")
    print(f"{test_name}")
    print(f"{'='*80}{Colors.ENDC}")

    if prompt_translation:
        print(f"{Colors.CYAN}Prompt (Translation): {prompt_translation}{Colors.ENDC}\n")

    print(f"{Colors.YELLOW}Generating from BASELINE model (original Mistral)...{Colors.ENDC}")
    baseline_output = generate_text(
        baseline_model, baseline_tokenizer, prompt, 
        max_new_tokens, temperature
    )

    print(f"{Colors.YELLOW}Generating from PRETRAINED model (after Korean training)...{Colors.ENDC}")
    pretrained_output = generate_text(
        pretrained_model, pretrained_tokenizer, prompt, 
        max_new_tokens, temperature
    )

    print(f"{Colors.YELLOW}Generating from FINETUNED model (after instruction tuning)...{Colors.ENDC}")
    finetuned_output = generate_text(
        finetuned_model, finetuned_tokenizer, prompt,
        max_new_tokens, temperature
    )

    # Display all three outputs
    print(f"\n{Colors.RED}┌─ BASELINE MODEL (Original Mistral) ───────────────────────────────┐{Colors.ENDC}")
    print(f"{Colors.RED}│{Colors.ENDC}")
    for line in baseline_output.split('\n'):
        print(f"{Colors.RED}│{Colors.ENDC} {line}")
    print(f"{Colors.RED}│{Colors.ENDC}")
    print(f"{Colors.RED}└────────────────────────────────────────────────────────────────────┘{Colors.ENDC}\n")

    print(f"{Colors.GREEN}┌─ PRETRAINED MODEL (After Korean Wikipedia) ───────────────────────┐{Colors.ENDC}")
    print(f"{Colors.GREEN}│{Colors.ENDC}")
    for line in pretrained_output.split('\n'):
        print(f"{Colors.GREEN}│{Colors.ENDC} {line}")
    print(f"{Colors.GREEN}│{Colors.ENDC}")
    print(f"{Colors.GREEN}└────────────────────────────────────────────────────────────────────┘{Colors.ENDC}\n")

    print(f"{Colors.CYAN}┌─ FINETUNED MODEL (After Instruction Tuning) ──────────────────────┐{Colors.ENDC}")
    print(f"{Colors.CYAN}│{Colors.ENDC}")
    for line in finetuned_output.split('\n'):
        print(f"{Colors.CYAN}│{Colors.ENDC} {line}")
    print(f"{Colors.CYAN}│{Colors.ENDC}")
    print(f"{Colors.CYAN}└────────────────────────────────────────────────────────────────────┘{Colors.ENDC}\n")

def parse_args():
    parser = argparse.ArgumentParser(
        description="对比韩语 Mistral 的三个阶段模型：基础模型 → 继续预训练 → 指令微调。"
                    "在同一批中韩英提示上并排生成，直观展示韩语能力的提升与英语能力的保留。",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--base_model", type=str, default="unsloth/mistral-7b-v0.3",
                        help="基础（未训练）模型名称")
    parser.add_argument("--pretrained_path", type=str, default="lora_model_pretrained",
                        help="继续预训练后保存的 LoRA 模型目录")
    parser.add_argument("--finetuned_path", type=str, default="lora_model",
                        help="指令微调后保存的最终 LoRA 模型目录")
    parser.add_argument("--max_seq_length", type=int, default=2048,
                        help="最大序列长度")
    parser.add_argument("--max_new_tokens", type=int, default=150,
                        help="每次生成的最大 token 数")
    parser.add_argument("--temperature", type=float, default=0.3,
                        help="采样温度（越低越确定）")
    return parser.parse_args()


def main():
    args = parse_args()

    print_section("🔬 KOREAN MISTRAL 3-WAY MODEL COMPARISON", Colors.HEADER)
    print(f"{Colors.YELLOW}This script compares three model stages:{Colors.ENDC}")
    print(f"  1. {Colors.RED}Baseline{Colors.ENDC} - Original Mistral (no Korean training)")
    print(f"  2. {Colors.GREEN}Pretrained{Colors.ENDC} - After Korean Wikipedia training")
    print(f"  3. {Colors.CYAN}Finetuned{Colors.ENDC} - After instruction tuning")
    print(f"\n{Colors.CYAN}Generation settings: temperature=0.3, do_sample=True (no repetition_penalty){Colors.ENDC}\n")

    # Load all three models
    print_section("📥 LOADING MODELS", Colors.BLUE)

    print(f"{Colors.YELLOW}Loading baseline model (original Mistral v0.3)...{Colors.ENDC}")
    baseline_model, baseline_tokenizer = load_baseline_model()
    print(f"{Colors.GREEN}✓ Baseline model loaded{Colors.ENDC}")

    print(f"\n{Colors.YELLOW}Loading pretrained model (after Korean pretraining)...{Colors.ENDC}")
    pretrained_model, pretrained_tokenizer = load_model("lora_model_pretrained")
    print(f"{Colors.GREEN}✓ Pretrained model loaded{Colors.ENDC}")

    print(f"\n{Colors.YELLOW}Loading finetuned model (after instruction tuning)...{Colors.ENDC}")
    finetuned_model, finetuned_tokenizer = load_model("lora_model")
    print(f"{Colors.GREEN}✓ Finetuned model loaded{Colors.ENDC}")

    # Define prompts
    wikipedia_prompt_korean = """위키피디아 기사
### 제목: {}

### 기사:
{}"""

    wikipedia_prompt_english = """Wikipedia Article
### Title: {}

### Article:
{}"""

    alpaca_prompt_korean = """다음은 작업을 설명하는 명령입니다. 요청을 적절하게 완료하는 응답을 작성하세요.

### 지침:
{}

### 응답:
{}"""

    alpaca_prompt_english = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{}

### Response:
{}"""

    print_section("🧪 RUNNING 3-WAY COMPARISONS", Colors.CYAN)

    # Test 1: Korean Wikipedia
    compare_on_prompt(
        baseline_model, baseline_tokenizer,
        pretrained_model, pretrained_tokenizer,
        finetuned_model, finetuned_tokenizer,
        wikipedia_prompt_korean.format("인공지능", ""),
        "Test 1: Korean Wikipedia - Artificial Intelligence (인공지능)",
        "Wikipedia Article / Title: Artificial Intelligence / Article:",
        max_new_tokens=args.max_new_tokens, temperature=args.temperature
    )

    # Test 2: English Wikipedia - Preservation Check
    compare_on_prompt(
        baseline_model, baseline_tokenizer,
        pretrained_model, pretrained_tokenizer,
        finetuned_model, finetuned_tokenizer,
        wikipedia_prompt_english.format("Artificial Intelligence", ""),
        "Test 2: English Wikipedia - Artificial Intelligence (Preservation Check)",
        None,
        max_new_tokens=args.max_new_tokens, temperature=args.temperature
    )

    # Test 3: Korean Instruction - Kimchi
    compare_on_prompt(
        baseline_model, baseline_tokenizer,
        pretrained_model, pretrained_tokenizer,
        finetuned_model, finetuned_tokenizer,
        alpaca_prompt_korean.format("한국의 전통 음식인 김치에 대해 설명하세요.", ""),
        "Test 3: Korean Instruction - Explain Kimchi",
        "Instruction: Explain about kimchi, a traditional Korean food. / Response:",
        max_new_tokens=args.max_new_tokens, temperature=args.temperature
    )

    # Test 4: Korean Instruction - Seoul
    compare_on_prompt(
        baseline_model, baseline_tokenizer,
        pretrained_model, pretrained_tokenizer,
        finetuned_model, finetuned_tokenizer,
        alpaca_prompt_korean.format("대한민국의 수도인 서울에 대해 간단히 소개해주세요.", ""),
        "Test 4: Korean Instruction - Introduce Seoul",
        "Instruction: Briefly introduce Seoul, the capital of South Korea. / Response:",
        max_new_tokens=args.max_new_tokens, temperature=args.temperature
    )

    # Test 5: English Instruction - Preservation Check
    compare_on_prompt(
        baseline_model, baseline_tokenizer,
        pretrained_model, pretrained_tokenizer,
        finetuned_model, finetuned_tokenizer,
        alpaca_prompt_english.format("Explain about Thanksgiving turkey, a traditional American food.", ""),
        "Test 5: English Instruction - Thanksgiving Turkey (Preservation Check)",
        None,
        max_new_tokens=args.max_new_tokens, temperature=args.temperature
    )

    print_section("📊 COMPARISON COMPLETE", Colors.GREEN)

    print(f"{Colors.CYAN}{'='*80}")
    print(f"💡 What to Look For:")
    print(f"{'='*80}{Colors.ENDC}")

    print(f"\n{Colors.RED}Baseline Model (Red boxes - Original Mistral):{Colors.ENDC}")
    print(f"  • Korean: Should be POOR - repetitive, nonsensical")
    print(f"  • English: Should be GOOD - this is the starting point")
    print(f"  • Shows what model knows BEFORE any Korean training")

    print(f"\n{Colors.GREEN}Pretrained Model (Green boxes - After Korean Wikipedia):{Colors.ENDC}")
    print(f"  • Korean: Should show IMPROVED fluency and vocabulary")
    print(f"  • Better Korean sentence structure")
    print(f"  • Weak instruction-following (only learned language, not how to follow instructions)")
    print(f"  • English: Should REMAIN strong (no catastrophic forgetting)")

    print(f"\n{Colors.CYAN}Finetuned Model (Cyan boxes - After Instruction Tuning):{Colors.ENDC}")
    print(f"  • Korean: Should be FLUENT with GOOD instruction-following")
    print(f"  • More structured and complete responses")
    print(f"  • Directly answers questions")
    print(f"  • English: Should REMAIN strong")

    print(f"\n{Colors.YELLOW}Key Progression to Observe:{Colors.ENDC}")
    print(f"  📊 Korean Quality:    {Colors.RED}Poor{Colors.ENDC} → {Colors.GREEN}Better{Colors.ENDC} → {Colors.CYAN}Best{Colors.ENDC}")
    print(f"  📊 Instruction:       {Colors.RED}Weak{Colors.ENDC} → {Colors.GREEN}Weak{Colors.ENDC} → {Colors.CYAN}Strong{Colors.ENDC}")
    print(f"  📊 English Quality:   {Colors.RED}Good{Colors.ENDC} → {Colors.GREEN}Good{Colors.ENDC} → {Colors.CYAN}Good{Colors.ENDC}")
    print(f"  📊 Repetition:        {Colors.RED}High{Colors.ENDC} → {Colors.GREEN}Medium{Colors.ENDC} → {Colors.CYAN}Low{Colors.ENDC}")

    print(f"\n{Colors.YELLOW}This demonstrates:{Colors.ENDC}")
    print(f"  ✓ Continued pretraining successfully teaches new language (Korean)")
    print(f"  ✓ Instruction tuning teaches how to follow instructions in the new language")
    print(f"  ✓ English capability is preserved throughout (no catastrophic forgetting)")
    print(f"  ✓ Both Wikipedia and Instruction tasks show English preservation")
    print(f"  ✓ Two-stage approach is necessary: language first, then instruction-following")
    print(f"\n{Colors.CYAN}💡 Note: Compare the English tests (Tests 2 & 5) across all three models.")
    print(f"All three should perform similarly well, proving no English degradation.{Colors.ENDC}")
    print()

if __name__ == "__main__":
    main()

continued-pretrain.py

# -*- 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,

        # 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,

        # 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")

evaluate_model.py

# -*- coding: utf-8 -*-
"""
Evaluation script for Korean Mistral continued-pretrained models
Loads saved LoRA adapters and evaluates on Korean and English tasks
"""

import os
import argparse

# 说明：unsloth / torch / transformers 等重型依赖在函数内按需导入，
# 这样 `python evaluate_model.py --help` 无需 GPU 环境即可查看参数。

# 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")

def load_model(model_path, max_seq_length=2048, dtype=None, load_in_4bit=True):
    """Load the saved LoRA model"""
    from unsloth import FastLanguageModel
    print_section(f"📥 LOADING MODEL FROM: {model_path}", Colors.BLUE)

    print(f"{Colors.YELLOW}Loading model and tokenizer...{Colors.ENDC}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name = model_path,
        max_seq_length = max_seq_length,
        dtype = dtype,
        load_in_4bit = load_in_4bit,
    )

    FastLanguageModel.for_inference(model)  # Enable native 2x faster inference
    print(f"{Colors.GREEN}✓ Model loaded successfully!{Colors.ENDC}")

    return model, tokenizer

def run_evaluation(model, tokenizer, max_new_tokens=150, 
                   temperature=0.7, top_p=0.9, use_sampling=False):
    """Run all evaluation tests"""
    from transformers import TextStreamer

    print_section("🧪 RUNNING EVALUATION TESTS", Colors.CYAN)
    print(f"{Colors.YELLOW}Generation Parameters:{Colors.ENDC}")
    print(f"  • max_new_tokens: {max_new_tokens}")
    if use_sampling:
        print(f"  • temperature: {temperature}")
        print(f"  • top_p: {top_p}")
        print(f"  • Sampling: Enabled")
    else:
        print(f"  • Sampling: Disabled (greedy decoding)")

    text_streamer = TextStreamer(tokenizer, skip_special_tokens=True)

    # Define prompts
    wikipedia_prompt_korean = """위키피디아 기사
### 제목: {}

### 기사:
{}"""

    wikipedia_prompt_english = """Wikipedia Article
### Title: {}

### Article:
{}"""

    alpaca_prompt_korean = """다음은 작업을 설명하는 명령입니다. 요청을 적절하게 완료하는 응답을 작성하세요.

### 지침:
{}

### 응답:
{}"""

    alpaca_prompt_english = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{}

### Response:
{}"""

    # Prepare generation kwargs
    gen_kwargs = {
        "max_new_tokens": max_new_tokens,
        "use_cache": True,
    }

    if use_sampling:
        gen_kwargs.update({
            "do_sample": True,
            "temperature": temperature,
            "top_p": top_p,
        })

    # Test 1: Korean Wikipedia Article
    print(f"\n{Colors.BOLD}{'='*70}")
    print(f"Test 1: Korean Wikipedia Article - Artificial Intelligence (인공지능)")
    print(f"{'='*70}{Colors.ENDC}")
    print(f"{Colors.CYAN}Prompt (Translation): Wikipedia Article / Title: Artificial Intelligence / Article:{Colors.ENDC}\n")

    test_prompt = wikipedia_prompt_korean.format("인공지능", "")
    inputs = tokenizer([test_prompt], return_tensors="pt").to("cuda")
    print(f"{Colors.GREEN}[KOREAN OUTPUT]{Colors.ENDC}")
    _ = model.generate(**inputs, streamer=text_streamer, **gen_kwargs)
    print(f"\n{Colors.BOLD}{'='*70}{Colors.ENDC}\n")

    # Test 2: English Wikipedia Article
    print(f"\n{Colors.BOLD}{'='*70}")
    print(f"Test 2: English Wikipedia Article - Artificial Intelligence")
    print(f"{'='*70}{Colors.ENDC}\n")

    test_prompt = wikipedia_prompt_english.format("Artificial Intelligence", "")
    inputs = tokenizer([test_prompt], return_tensors="pt").to("cuda")
    print(f"{Colors.GREEN}[ENGLISH OUTPUT]{Colors.ENDC}")
    _ = model.generate(**inputs, streamer=text_streamer, **gen_kwargs)
    print(f"\n{Colors.BOLD}{'='*70}{Colors.ENDC}\n")

    # Test 3: Korean Instruction (Kimchi)
    print(f"\n{Colors.BOLD}{'='*70}")
    print(f"Test 3: Korean Instruction - Explain about Kimchi")
    print(f"{'='*70}{Colors.ENDC}")
    print(f"{Colors.CYAN}Prompt (Translation): Instruction: Explain about kimchi, a traditional Korean food. / Response:{Colors.ENDC}\n")

    test_prompt = alpaca_prompt_korean.format("한국의 전통 음식인 김치에 대해 설명하세요.", "")
    inputs = tokenizer([test_prompt], return_tensors="pt").to("cuda")
    print(f"{Colors.GREEN}[KOREAN OUTPUT]{Colors.ENDC}")
    _ = model.generate(**inputs, streamer=text_streamer, **gen_kwargs)
    print(f"\n{Colors.BOLD}{'='*70}{Colors.ENDC}\n")

    # Test 4: English Instruction (Thanksgiving)
    print(f"\n{Colors.BOLD}{'='*70}")
    print(f"Test 4: English Instruction - Explain about Thanksgiving Turkey")
    print(f"{'='*70}{Colors.ENDC}\n")

    test_prompt = alpaca_prompt_english.format("Explain about Thanksgiving turkey, a traditional American food.", "")
    inputs = tokenizer([test_prompt], return_tensors="pt").to("cuda")
    print(f"{Colors.GREEN}[ENGLISH OUTPUT]{Colors.ENDC}")
    _ = model.generate(**inputs, streamer=text_streamer, **gen_kwargs)
    print(f"\n{Colors.BOLD}{'='*70}{Colors.ENDC}\n")

    # Additional Korean tests
    print(f"\n{Colors.BOLD}{'='*70}")
    print(f"Test 5: Korean Instruction - Explain about Seoul")
    print(f"{'='*70}{Colors.ENDC}")
    print(f"{Colors.CYAN}Prompt (Translation): Instruction: Briefly introduce Seoul, the capital of South Korea. / Response:{Colors.ENDC}\n")

    test_prompt = alpaca_prompt_korean.format("대한민국의 수도인 서울에 대해 간단히 소개해주세요.", "")
    inputs = tokenizer([test_prompt], return_tensors="pt").to("cuda")
    print(f"{Colors.GREEN}[KOREAN OUTPUT]{Colors.ENDC}")
    _ = model.generate(**inputs, streamer=text_streamer, **gen_kwargs)
    print(f"\n{Colors.BOLD}{'='*70}{Colors.ENDC}\n")

    # Test 6: Korean Instruction (K-pop)
    print(f"\n{Colors.BOLD}{'='*70}")
    print(f"Test 6: Korean Instruction - Explain about K-pop")
    print(f"{'='*70}{Colors.ENDC}")
    print(f"{Colors.CYAN}Prompt (Translation): Instruction: Explain what K-pop is. / Response:{Colors.ENDC}\n")

    test_prompt = alpaca_prompt_korean.format("K-pop이 무엇인지 설명해주세요.", "")
    inputs = tokenizer([test_prompt], return_tensors="pt").to("cuda")
    print(f"{Colors.GREEN}[KOREAN OUTPUT]{Colors.ENDC}")
    _ = model.generate(**inputs, streamer=text_streamer, **gen_kwargs)
    print(f"\n{Colors.BOLD}{'='*70}{Colors.ENDC}\n")

    print_section("✅ EVALUATION COMPLETE", Colors.GREEN)

def main():
    parser = argparse.ArgumentParser(description="Evaluate Korean Mistral LoRA models")
    parser.add_argument(
        "--model_path",
        type=str,
        default="lora_model",
        help="Path to the saved LoRA model (default: lora_model)"
    )
    parser.add_argument(
        "--pretrained",
        action="store_true",
        help="Load the pretrained model (before SFT) instead of final model"
    )
    parser.add_argument(
        "--max_seq_length",
        type=int,
        default=2048,
        help="Maximum sequence length (default: 2048)"
    )
    parser.add_argument(
        "--load_in_4bit",
        action="store_true",
        default=True,
        help="Load model in 4-bit quantization (default: True)"
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=150,
        help="Maximum number of tokens to generate (default: 150)"
    )
    parser.add_argument(
        "--use_sampling",
        action="store_true",
        help="Use sampling instead of greedy decoding"
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.7,
        help="Sampling temperature (default: 0.7, only used with --use_sampling)"
    )
    parser.add_argument(
        "--top_p",
        type=float,
        default=0.9,
        help="Top-p sampling parameter (default: 0.9, only used with --use_sampling)"
    )

    args = parser.parse_args()

    # Determine model path
    if args.pretrained:
        model_path = "lora_model_pretrained"
        print(f"{Colors.YELLOW}Loading PRETRAINED model (before instruction finetuning){Colors.ENDC}")
    else:
        model_path = args.model_path
        print(f"{Colors.YELLOW}Loading FINETUNED model (after instruction finetuning){Colors.ENDC}")

    # Check if model exists
    if not os.path.exists(model_path):
        print(f"{Colors.RED}Error: Model path '{model_path}' does not exist!{Colors.ENDC}")
        print(f"{Colors.YELLOW}Make sure you've run the training script first.{Colors.ENDC}")
        return

    print_section("🚀 KOREAN MISTRAL MODEL EVALUATION", Colors.HEADER)

    # Load model
    model, tokenizer = load_model(
        model_path=model_path,
        max_seq_length=args.max_seq_length,
        load_in_4bit=args.load_in_4bit
    )

    # Display GPU info
    import torch
    print_section("💾 GPU MEMORY STATS", Colors.YELLOW)
    gpu_stats = torch.cuda.get_device_properties(0)
    reserved_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}")
    print(f"Max memory: {max_memory} GB")
    print(f"Reserved memory: {reserved_memory} GB")

    # Run evaluation
    run_evaluation(
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        use_sampling=args.use_sampling
    )

    print(f"\n{Colors.CYAN}{'='*70}")
    print(f"💡 Tips:")
    print(f"{'='*70}{Colors.ENDC}")
    print(f"• Compare pretrained vs finetuned: Run with --pretrained flag")
    print(f"• Adjust generation: Use --max_new_tokens")
    print(f"• Enable sampling: Use --use_sampling --temperature 0.7 --top_p 0.9")
    print(f"• Example: python evaluate_model.py --pretrained --max_new_tokens 300")
    print()

if __name__ == "__main__":
    main()