network_admin_llm_train.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
'''
Network Admin LLM - QLoRA Fine-tuning Script
=============================================
Base Model: microsoft/Phi-4-mini-instruct
Method: QLoRA SFT (4-bit quantization + LoRA)
Datasets: NetEval + Telecom Intent Config

Run locally with GPU:
    pip install transformers trl peft bitsandbytes accelerate datasets trackio
    python network_admin_llm_train.py

Or on Google Colab:
    !pip install transformers trl peft bitsandbytes accelerate datasets trackio
    %cd /content
    !python network_admin_llm_train.py

Author: Network Admin LLM Project
'''

import os
import sys
import torch
from datetime import datetime

# ============== CONFIGURATION ==============
# 請修改以下設定
MODEL_NAME = 'microsoft/Phi-4-mini-instruct'
HF_USERNAME = 'YOUR_HF_USERNAME'  # 改成你的 HuggingFace 用戶名
HF_TOKEN = os.environ.get('HF_TOKEN', 'YOUR_HF_TOKEN')  # HF token for upload

# 訓練超參數
TRAINING_CONFIG = {
    'learning_rate': 2e-4,          # LoRA 需要較高學習率
    'num_epochs': 3,
    'batch_size': 4,
    'gradient_accumulation': 4,      # effective batch = 16
    'max_seq_length': 2048,
    'lora_r': 16,
    'lora_alpha': 32,
    'lora_dropout': 0.05,
    'warmup_ratio': 0.1,
}

OUTPUT_DIR = f'{HF_USERNAME}/network-admin-phi4-mini'
# ===========================================

def print_section(title):
    print(f'\n{"="*60}')
    print(f' {title}')
    print('='*60)

def install_dependencies():
    '''檢查並安裝依賴'''
    print_section('CHECKING DEPENDENCIES')
    
    required = ['transformers', 'trl', 'peft', 'bitsandbytes', 'accelerate', 'datasets', 'trackio']
    missing = []
    
    for pkg in required:
        try:
            __import__(pkg.replace('-', '_'))
            print(f'✅ {pkg}')
        except ImportError:
            missing.append(pkg)
            print(f'❌ {pkg} - 需要安裝')
    
    if missing:
        print(f'\n請運行: pip install {" ".join(missing)}')
        return False
    return True

def load_and_prepare_datasets():
    '''載入並轉換數據集'''
    from datasets import load_dataset, concatenate_datasets
    
    print_section('LOADING DATASETS')
    
    # 1. 載入 NetEval 考試題庫
    print('📚 載入 NetEval 考試題庫...')
    neteval_dataset = load_dataset('NASP/neteval-exam', split='train')
    print(f'   NetEval: {len(neteval_dataset)} 題')
    
    def convert_neteval(example):
        '''將 Q&A 格式轉換為對話格式'''
        question = example['Question']
        options = f'\nA. {example.get("A", "")}\nB. {example.get("B", "")}\nC. {example.get("C", "")}\nD. {example.get("D", "")}'
        
        answer = f'正確答案是: {example["Answer"]}'
        if example.get('Explanation'):
            answer += f'\n\n📖 解說: {example["Explanation"]}'
        
        return {
            'messages': [
                {'role': 'system', 'content': '你是一位網路管理專家。請回答關於網路、安全、路由、交換機、VLAN、防火牆等IT基礎設施的問題。'},
                {'role': 'user', 'content': f'{question}{options}'},
                {'role': 'assistant', 'content': answer}
            ]
        }
    
    neteval_converted = neteval_dataset.map(
        convert_neteval,
        remove_columns=neteval_dataset.column_names,
        desc='轉換 NetEval 格式'
    )
    
    # 2. 載入電信意圖配置數據集
    print('📚 載入電信意圖配置數據集...')
    telecom_dataset = load_dataset('nraptisss/telecom-intent-config-sft-10k', split='train')
    print(f'   Telecom: {len(telecom_dataset)} 條')
    
    telecom_messages = telecom_dataset.map(
        lambda x: {'messages': x['messages']},
        remove_columns=[c for c in telecom_dataset.column_names if c != 'messages']
    )
    
    # 3. 合併數據集
    print('🔄 合併數據集...')
    combined = concatenate_datasets([neteval_converted, telecom_messages])
    split_data = combined.train_test_split(test_size=0.1, seed=42)
    
    train_ds = split_data['train']
    eval_ds = split_data['test']
    
    print(f'\n📊 數據集統計:')
    print(f'   訓練集: {len(train_ds)} 條')
    print(f'   驗證集: {len(eval_ds)} 條')
    print(f'   總計: {len(combined)} 條')
    
    return train_ds, eval_ds

def setup_model_and_tokenizer():
    '''設置模型和 tokenizer'''
    from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
    from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
    
    print_section('LOADING MODEL')
    print(f'🤖 模型: {MODEL_NAME}')
    
    # Tokenizer
    print('\n📝 載入 Tokenizer...')
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = 'right'
    print(f'   Vocab size: {len(tokenizer):,}')
    
    # QLoRA 配置 (4-bit)
    print('\n⚡ 配置 QLoRA (4-bit)...')
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type='nf4',          # Normalized Float4
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,      # 嵌套量化
    )
    
    # 載入模型
    print('📥 載入模型 (4-bit)...')
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map='auto',
        trust_remote_code=True,
    )
    
    # 準備 kbit 訓練
    model = prepare_model_for_kbit_training(model)
    print('✅ 模型準備完成')
    
    # LoRA 配置
    print('\n🔧 配置 LoRA...')
    lora_config = LoraConfig(
        r=TRAINING_CONFIG['lora_r'],
        lora_alpha=TRAINING_CONFIG['lora_alpha'],
        lora_dropout=TRAINING_CONFIG['lora_dropout'],
        bias='none',
        task_type='CAUSAL_LM',
        target_modules=[
            'q_proj', 'k_proj', 'v_proj', 'o_proj',  # Attention
            'gate_proj', 'up_proj', 'down_proj',      # MLP
        ],
        modules_to_save=['lm_head', 'embed_tokens'],
    )
    
    # 應用 LoRA
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    
    return model, tokenizer, lora_config

def setup_trainer(model, tokenizer, train_ds, eval_ds, lora_config):
    '''設置訓練器'''
    from trl import SFTTrainer, SFTConfig
    
    print_section('CONFIGURING TRAINER')
    
    # 生成運行名稱
    run_name = f'phi4-netadmin-{datetime.now().strftime("%m%d-%H%M")}'
    
    # 嘗試初始化 trackio
    try:
        import trackio
        trackio.init(project='network-admin-llm', experiment='qlora-sft', run_name=run_name)
        print('✅ Trackio 初始化成功')
        report_to = ['trackio']
    except Exception as e:
        print(f'⚠️ Trackio 初始化失敗: {e}')
        report_to = ['none']
    
    # SFT 配置
    training_args = SFTConfig(
        # 學習率
        learning_rate=TRAINING_CONFIG['learning_rate'],
        lr_scheduler_type='cosine',
        warmup_ratio=TRAINING_CONFIG['warmup_ratio'],
        
        # 訓練
        num_train_epochs=TRAINING_CONFIG['num_epochs'],
        per_device_train_batch_size=TRAINING_CONFIG['batch_size'],
        gradient_accumulation_steps=TRAINING_CONFIG['gradient_accumulation'],
        max_seq_length=TRAINING_CONFIG['max_seq_length'],
        
        # 記憶體優化
        gradient_checkpointing=True,
        bf16=True,
        fp16=False,
        
        # 輸出
        output_dir='./output',
        logging_steps=10,
        save_steps=500,
        save_total_limit=2,
        evaluation_strategy='steps',
        eval_steps=500,
        
        # Hub 上傳
        push_to_hub=True,
        hub_model_id=OUTPUT_DIR,
        hub_strategy='checkpoint',
        
        # 監控
        report_to=report_to,
        logging_strategy='steps',
        logging_first_step=True,
        
        # 雜項
        remove_unused_columns=False,
        dataloader_num_workers=4,
        seed=42,
    )
    
    # 創建 trainer
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        processing_class=tokenizer,
        peft_config=lora_config,
    )
    
    return trainer, run_name

def train_model(trainer):
    '''執行訓練'''
    print_section('STARTING TRAINING')
    print('🚀 開始訓練...')
    print('   (按 Ctrl+C 可隨時中斷)')
    print()
    
    try:
        trainer.train()
        print('\n✅ 訓練完成!')
        return True
    except KeyboardInterrupt:
        print('\n⚠️ 訓練被用戶中斷')
        return False
    except Exception as e:
        print(f'\n❌ 訓練失敗: {e}')
        raise

def save_and_upload(trainer):
    '''保存並上傳模型'''
    print_section('SAVING & UPLOADING')
    
    try:
        print('📤 上傳模型到 HuggingFace Hub...')
        trainer.push_to_hub()
        print(f'\n✅ 模型已上傳!')
        print(f'🔗 連結: https://huggingface.co/{OUTPUT_DIR}')
    except Exception as e:
        print(f'\n⚠️ 上傳失敗: {e}')
        print('模型已保存在 ./output 目錄')

def main():
    '''主函數'''
    print('''
    ╔═══════════════════════════════════════════════════════════╗
    ║       Network Admin LLM - QLoRA Fine-tuning              ║
    ║       Base: microsoft/Phi-4-mini-instruct                ║
    ╚═══════════════════════════════════════════════════════════╝
    ''')
    
    # 檢查 GPU
    print(f'🖥️  GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else "無 GPU"}')
    if torch.cuda.is_available():
        print(f'   Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB')
    
    # 安裝依賴
    if not install_dependencies():
        sys.exit(1)
    
    # 載入數據
    train_ds, eval_ds = load_and_prepare_datasets()
    
    # 設置模型
    model, tokenizer, lora_config = setup_model_and_tokenizer()
    
    # 設置 trainer
    trainer, run_name = setup_trainer(model, tokenizer, train_ds, eval_ds, lora_config)
    
    # 訓練
    success = train_model(trainer)
    
    # 保存
    if success:
        save_and_upload(trainer)
    
    print_section('DONE')
    print(f'Run name: {run_name}')

if __name__ == '__main__':
    main()