RLHF: Reinforcement Learning from Human Feedback

from transformers import AutoModelForCausalLM, AutoTokenizer from trl import ( SFTTrainer, RewardTrainer, PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead ) from datasets import load_dataset import torch class RLHFPipeline: """ Pipeline completo de RLHF em 3 etapas: 1. Supervised Fine-Tuning (SFT) 2. Reward Model Training 3. PPO Optimization """ def __init__(self, base_model="meta-llama/Llama-2-7b-hf"): self.base_model = base_model self.tokenizer = AutoTokenizer.from_pretrained(base_model) self.tokenizer.pad_token = self.tokenizer.eos_token # ===================== # ETAPA 1: SFT # ===================== def supervised_fine_tuning(self, dataset_name="OpenAssistant/oasst1"): """ Treina modelo em conversas de alta qualidade escritas por humanos """ model = AutoModelForCausalLM.from_pretrained( self.base_model, torch_dtype=torch.bfloat16, device_map="auto" ) # Dataset de conversas humanas dataset = load_dataset(dataset_name, split="train") # Formata conversas def format_conversation(example): return { "text": f"Human: {example['prompt']}\n\nAssistant: {example['response']}" } dataset = dataset.map(format_conversation) # Trainer SFT trainer = SFTTrainer( model=model, train_dataset=dataset, dataset_text_field="text", max_seq_length=2048, tokenizer=self.tokenizer, ) trainer.train() trainer.save_model("./sft_model") return "./sft_model" # ===================== # ETAPA 2: REWARD MODEL # ===================== def train_reward_model(self, sft_model_path, comparison_dataset): """ Treina modelo para prever preferências humanas Dataset format: - prompt: pergunta do usuário - chosen: resposta preferida - rejected: resposta rejeitada """ from trl import RewardTrainer, RewardConfig from transformers import AutoModelForSequenceClassification # Modelo de reward reward_model = AutoModelForSequenceClassification.from_pretrained( sft_model_path, num_labels=1, torch_dtype=torch.bfloat16 ) # Dataset de comparações humanas dataset = load_dataset(comparison_dataset, split="train") # Config de treinamento config = RewardConfig( output_dir="./reward_model", per_device_train_batch_size=4, num_train_epochs=1, learning_rate=1e-5, ) trainer = RewardTrainer( model=reward_model, tokenizer=self.tokenizer, args=config, train_dataset=dataset, ) trainer.train() trainer.save_model("./reward_model") return "./reward_model" # ===================== # ETAPA 3: PPO TRAINING # ===================== def ppo_training(self, sft_model_path, reward_model_path): """ Otimiza modelo via PPO para maximizar reward enquanto mantém coerência (KL penalty) """ # Modelo com value head para PPO model = AutoModelForCausalLMWithValueHead.from_pretrained( sft_model_path, torch_dtype=torch.bfloat16, device_map="auto" ) # Modelo de referência (frozen) ref_model = AutoModelForCausalLMWithValueHead.from_pretrained( sft_model_path, torch_dtype=torch.bfloat16, device_map="auto" ) # Reward model reward_model = AutoModelForSequenceClassification.from_pretrained( reward_model_path, torch_dtype=torch.bfloat16, device_map="auto" ) # Configuração PPO ppo_config = PPOConfig( model_name=sft_model_path, learning_rate=1e-6, batch_size=16, mini_batch_size=4, gradient_accumulation_steps=4, ppo_epochs=4, kl_penalty="kl", # Penalidade KL init_kl_coef=0.2, # β na fórmula ) # PPO Trainer ppo_trainer = PPOTrainer( config=ppo_config, model=model, ref_model=ref_model, tokenizer=self.tokenizer, ) # Dataset de prompts prompts_dataset = load_dataset("Anthropic/hh-rlhf", split="train") # Loop de treinamento PPO for batch in prompts_dataset: # Gera respostas query_tensors = self.tokenizer( batch["prompt"], return_tensors="pt", padding=True ).input_ids response_tensors = ppo_trainer.generate( query_tensors, max_new_tokens=256, do_sample=True, temperature=0.7 ) # Calcula rewards texts = self.tokenizer.batch_decode(response_tensors) rewards = self._compute_rewards(reward_model, texts) # Step PPO stats = ppo_trainer.step( query_tensors, response_tensors, rewards ) print(f"Mean reward: {stats['ppo/mean_reward']:.4f}") print(f"KL divergence: {stats['ppo/kl']:.4f}") ppo_trainer.save_model("./rlhf_model") return "./rlhf_model" def _compute_rewards(self, reward_model, texts): """Calcula reward para cada resposta""" inputs = self.tokenizer( texts, return_tensors="pt", padding=True, truncation=True ) with torch.no_grad(): outputs = reward_model(**inputs) rewards = outputs.logits.squeeze() return rewards.tolist() # Constitutional AI (Anthropic) class ConstitutionalAI: """ RLHF com princípios constitucionais Modelo critica próprias respostas """ PRINCIPLES = [ "Seja útil, inofensivo e honesto", "Não ajude com atividades ilegais", "Admita quando não souber algo", "Respeite privacidade e segurança", "Evite viés e discriminação", ] def generate_critique(self, model, response, principle): """Gera crítica baseada em princípio""" prompt = f"""Avalie esta resposta segundo o princípio: {principle} Resposta: {response} Crítica:""" return model.generate(prompt, max_tokens=200) def revise_response(self, model, response, critique): """Revisa resposta baseada na crítica""" prompt = f"""Revise esta resposta considerando a crítica: Original: {response} Crítica: {critique} Revisão:""" return model.generate(prompt, max_tokens=500) # Uso completo if __name__ == "__main__": # Pipeline RLHF completo pipeline = RLHFPipeline(base_model="meta-llama/Llama-2-7b-hf") # Etapa 1: SFT print("=== Supervised Fine-Tuning ===") sft_path = pipeline.supervised_fine_tuning() # Etapa 2: Reward Model print("=== Training Reward Model ===") reward_path = pipeline.train_reward_model( sft_path, "Anthropic/hh-rlhf" ) # Etapa 3: PPO print("=== PPO Training ===") final_path = pipeline.ppo_training(sft_path, reward_path) print(f"Modelo RLHF salvo em: {final_path}")