定制模型训练：打造企业专属AI大脑

LoRA微调实战

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, TaskType, get_peft_model
import torch
from datasets import Dataset

class CustomModelTrainer:
    """企业定制模型训练器"""
    
    def __init__(self, base_model_path: str, output_dir: str):
        self.base_model_path = base_model_path
        self.output_dir = output_dir
        
        # 加载基础模型和分词器
        self.model = AutoModelForCausalLM.from_pretrained(
            base_model_path,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        
        self.tokenizer = AutoTokenizer.from_pretrained(base_model_path)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        
    def prepare_lora_model(self, lora_config: dict):
        """配置LoRA模型"""
        peft_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=lora_config.get('r', 16),
            lora_alpha=lora_config.get('alpha', 32),
            lora_dropout=lora_config.get('dropout', 0.1),
            target_modules=[
                "q_proj", "v_proj", "k_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj"
            ],
            inference_mode=False
        )
        
        self.model = get_peft_model(self.model, peft_config)
        self.model.print_trainable_parameters()
        
    def prepare_dataset(self, data_path: str):
        """准备训练数据集"""
        # 加载数据
        with open(data_path, 'r', encoding='utf-8') as f:
            raw_data = json.load(f)
        
        # 格式化为模型输入
        formatted_data = []
        for item in raw_data:
            text = self.format_instruction(
                item['instruction'],
                item['input'],
                item['output']
            )
            formatted_data.append({'text': text})
        
        # 转换为Dataset
        dataset = Dataset.from_list(formatted_data)
        
        # Tokenize
        def tokenize_function(examples):
            return self.tokenizer(
                examples['text'],
                truncation=True,
                padding=True,
                max_length=2048
            )
        
        tokenized_dataset = dataset.map(
            tokenize_function,
            batched=True,
            remove_columns=['text']
        )
        
        return tokenized_dataset
    
    def train(self, train_dataset, eval_dataset=None):
        """执行训练"""
        training_args = TrainingArguments(
            output_dir=self.output_dir,
            num_train_epochs=3,
            per_device_train_batch_size=4,
            gradient_accumulation_steps=4,
            warmup_steps=100,
            learning_rate=5e-4,
            fp16=True,
            logging_steps=10,
            evaluation_strategy="steps" if eval_dataset else "no",
            eval_steps=100 if eval_dataset else None,
            save_strategy="steps",
            save_steps=200,
            save_total_limit=3,
            load_best_model_at_end=True if eval_dataset else False,
            report_to="tensorboard",
            gradient_checkpointing=True,
        )
        
        # 数据整理器
        data_collator = DataCollatorForLanguageModeling(
            tokenizer=self.tokenizer,
            mlm=False
        )
        
        # 创建训练器
        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            data_collator=data_collator,
            callbacks=[
                EarlyStoppingCallback(early_stopping_patience=3),
                CustomLoggingCallback()
            ]
        )
        
        # 开始训练
        trainer.train()
        
        # 保存模型
        trainer.save_model()
        self.model.save_pretrained(self.output_dir)
        
    def evaluate_model(self, test_dataset):
        """评估模型效果"""
        results = {
            'perplexity': self.calculate_perplexity(test_dataset),
            'domain_accuracy': self.test_domain_knowledge(),
            'safety_score': self.safety_evaluation(),
            'inference_speed': self.benchmark_speed()
        }
        
        return results