DPO: Direct Preference Optimization

from transformers import AutoModelForCausalLM, AutoTokenizer from trl import DPOTrainer, DPOConfig from datasets import load_dataset, Dataset import torch class DPOPipeline: """ Pipeline DPO simplificado: 1. Carrega modelo SFT 2. Prepara dataset de preferências 3. Treina com DPO loss Sem reward model, sem PPO, sem complexidade! """ def __init__(self, base_model="meta-llama/Llama-2-7b-chat-hf"): self.base_model = base_model self.tokenizer = AutoTokenizer.from_pretrained(base_model) self.tokenizer.pad_token = self.tokenizer.eos_token def prepare_preference_dataset(self, dataset_name="Anthropic/hh-rlhf"): """ Prepara dataset no formato DPO: - prompt: input do usuário - chosen: resposta preferida - rejected: resposta rejeitada """ dataset = load_dataset(dataset_name, split="train") def format_example(example): # Extrai prompt, chosen e rejected # Formato varia por dataset return { "prompt": example["prompt"], "chosen": example["chosen"], "rejected": example["rejected"] } return dataset.map(format_example) def train_dpo( self, dataset, output_dir="./dpo_model", beta=0.1, learning_rate=5e-7, num_epochs=1 ): """ Treina modelo com DPO Args: dataset: Dataset com prompt, chosen, rejected beta: Peso do KL penalty (mais alto = mais conservador) learning_rate: LR baixo para fine-tuning num_epochs: Geralmente 1-3 épocas suficientes """ # Carrega modelo model = AutoModelForCausalLM.from_pretrained( self.base_model, torch_dtype=torch.bfloat16, device_map="auto" ) # Modelo de referência (frozen) ref_model = AutoModelForCausalLM.from_pretrained( self.base_model, torch_dtype=torch.bfloat16, device_map="auto" ) # Configuração DPO dpo_config = DPOConfig( output_dir=output_dir, beta=beta, # β na fórmula learning_rate=learning_rate, num_train_epochs=num_epochs, per_device_train_batch_size=4, gradient_accumulation_steps=4, warmup_ratio=0.1, logging_steps=10, save_steps=100, bf16=True, # Otimizações de memória gradient_checkpointing=True, optim="adamw_8bit", ) # Trainer DPO trainer = DPOTrainer( model=model, ref_model=ref_model, args=dpo_config, train_dataset=dataset, tokenizer=self.tokenizer, ) # Treina! trainer.train() # Salva modelo trainer.save_model(output_dir) return output_dir # Variantes de DPO class IPOTraining: """ IPO: Identity Preference Optimization Remove sigmoid, mais estável para preferências extremas """ def compute_ipo_loss(self, policy_chosen_logps, policy_rejected_logps, ref_chosen_logps, ref_rejected_logps, beta=0.1): """ IPO Loss - sem sigmoid """ chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps) rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps) # Squared difference ao invés de log-sigmoid loss = (chosen_rewards - rejected_rewards - 1) ** 2 return loss.mean() class KTOTraining: """ KTO: Kahneman-Tversky Optimization Funciona com feedback binário (bom/ruim) Não precisa de pares comparativos! """ def compute_kto_loss(self, policy_logps, ref_logps, is_chosen, beta=0.1): """ KTO Loss - feedback binário Args: is_chosen: True se resposta é boa, False se ruim """ logratios = policy_logps - ref_logps if is_chosen: # Maximiza para respostas boas loss = 1 - torch.sigmoid(beta * logratios) else: # Minimiza para respostas ruins loss = 1 - torch.sigmoid(-beta * logratios) return loss.mean() class ORPOTraining: """ ORPO: Odds Ratio Preference Optimization Combina SFT + alinhamento em uma única etapa Mais eficiente que DPO puro """ def compute_orpo_loss(self, policy_chosen_logps, policy_rejected_logps, labels, lambda_orpo=1.0): """ ORPO = SFT Loss + λ * Odds Ratio Loss """ # SFT loss (NLL) sft_loss = -policy_chosen_logps.mean() # Odds ratio chosen_odds = torch.exp(policy_chosen_logps) rejected_odds = torch.exp(policy_rejected_logps) odds_ratio = chosen_odds / rejected_odds orpo_loss = -torch.log(odds_ratio / (1 + odds_ratio)) # Combinação total_loss = sft_loss + lambda_orpo * orpo_loss.mean() return total_loss # Criando dataset de preferências customizado class PreferenceDatasetBuilder: """ Constrói dataset de preferências para DPO a partir de anotações humanas ou sintéticas """ def __init__(self, model_for_generation=None): self.generator = model_for_generation def create_from_human_annotations( self, prompts, responses_a, responses_b, preferences # 'a' ou 'b' ): """ Cria dataset a partir de comparações humanas """ data = [] for prompt, resp_a, resp_b, pref in zip( prompts, responses_a, responses_b, preferences ): if pref == 'a': chosen, rejected = resp_a, resp_b else: chosen, rejected = resp_b, resp_a data.append({ "prompt": prompt, "chosen": chosen, "rejected": rejected }) return Dataset.from_list(data) def create_synthetic_preferences( self, prompts, good_model, bad_model ): """ Gera preferências sintéticas usando modelo bom vs modelo ruim (RLAIF style) """ data = [] for prompt in prompts: # Gera respostas de ambos modelos chosen = good_model.generate(prompt) rejected = bad_model.generate(prompt) data.append({ "prompt": prompt, "chosen": chosen, "rejected": rejected }) return Dataset.from_list(data) # Uso completo if __name__ == "__main__": # Pipeline DPO simples pipeline = DPOPipeline( base_model="meta-llama/Llama-2-7b-chat-hf" ) # Prepara dataset print("=== Preparando Dataset ===") dataset = pipeline.prepare_preference_dataset( "Anthropic/hh-rlhf" ) print(f"Dataset size: {len(dataset)} exemplos") # Treina com DPO print("\n=== Treinando DPO ===") model_path = pipeline.train_dpo( dataset, output_dir="./dpo_llama", beta=0.1, learning_rate=5e-7, num_epochs=1 ) print(f"\nModelo DPO salvo em: {model_path}") # Comparação: DPO vs RLHF print("\n=== Comparação ===") print("RLHF: 3 modelos, PPO, 100% custo") print("DPO: 1 modelo, supervised, 50% custo") print("Qualidade: ~99% equivalente")