import torch import torch.nn as nn from transformers import ( GPT2LMHeadModel, AutoTokenizer, RobertaForSequenceClassification, GPT2Config, RobertaConfig, DataCollatorForLanguageModeling ) from datasets import load_dataset from torch.utils.data import DataLoader from accelerate import Accelerator import logging # --- Configuration --- # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Constants GPT2_MODEL_NAME = "gpt2" ROBERTA_MODEL_NAME = "roberta-base" SEQUENCE_LENGTH = 128 BATCH_SIZE = 8 NUM_EPOCHS = 3 LEARNING_RATE_G = 1e-5 # Lower LR for generation models LEARNING_RATE_D = 5e-5 # Higher LR for classification models D_STEPS = 1 # Number of discriminator updates per generator update G_STEPS = 1 # Number of generator updates per batch # --- 1. Discriminator Wrapper Class --- # We wrap RoBERTa to make it function as a binary classifier (0: Fake, 1: Real) class Discriminator(nn.Module): def __init__(self, model_name): super().__init__() # RoBERTa is loaded for sequence classification with 2 labels (real/fake) self.roberta = RobertaForSequenceClassification.from_pretrained(model_name, num_labels=2) def forward(self, input_ids, attention_mask=None, labels=None): # The RoBERTa model outputs a SequenceClassifierOutput output = self.roberta( input_ids=input_ids, attention_mask=attention_mask, labels=labels ) # We only need the logits for the GAN loss calculation return output.logits # --- 2. Tokenizer Initialization (Kept Global for Data Preprocessing) --- # Initialize the tokenizer tokenizer = AutoTokenizer.from_pretrained(GPT2_MODEL_NAME) # GPT-2 does not have a native padding token, so we set the EOS token as the pad token # This is crucial for batching and RoBERTa's input structure if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token # --- 3. Data Preprocessing --- # This function requires the tokenizer to be available globally def preprocess_function(examples): # Tokenize the dataset return tokenizer(examples["text"], max_length=SEQUENCE_LENGTH, truncation=True, padding="max_length") def load_and_prepare_data(): # Load a dataset of real text (e.g., IMDB reviews) raw_datasets = load_dataset("imdb", split="train[:5%]") # Select only the 'text' column for language modeling processed_datasets = raw_datasets.map( preprocess_function, batched=True, remove_columns=raw_datasets.column_names, ) # Convert to PyTorch tensors and prepare for DataLoader processed_datasets.set_format(type="torch", columns=["input_ids", "attention_mask"]) # Simple data collator for padding data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) # Create DataLoader dataloader = DataLoader( processed_datasets, shuffle=True, collate_fn=data_collator, # The data collator will handle batching and masking batch_size=BATCH_SIZE ) return dataloader # --- 4. Adversarial Training Function --- def train_gan(): # Initialize Generator and Discriminator models LOCALLY to ensure correct scope logger.info(f"Loading Generator ({GPT2_MODEL_NAME}) and Discriminator ({ROBERTA_MODEL_NAME}) inside train_gan...") generator = GPT2LMHeadModel.from_pretrained(GPT2_MODEL_NAME) discriminator = Discriminator(ROBERTA_MODEL_NAME) # Initialize Accelerator for mixed-precision and distributed training handling accelerator = Accelerator() dataloader = load_and_prepare_data() # Define optimizers # These references now correctly point to the locally defined 'generator' and 'discriminator' optimizer_g = torch.optim.AdamW(generator.parameters(), lr=LEARNING_RATE_G) optimizer_d = torch.optim.AdamW(discriminator.parameters(), lr=LEARNING_RATE_D) # Move models and optimizers to the appropriate device generator, optimizer_g, discriminator, optimizer_d, dataloader = accelerator.prepare( generator, optimizer_g, discriminator, optimizer_d, dataloader ) # Define Loss Function: Binary Cross-Entropy with Logits # Since RoBERTa is outputting logits (unscaled scores), BCEWithLogitsLoss is the correct, stable choice. loss_fn = nn.BCEWithLogitsLoss() logger.info("Starting adversarial training loop...") # Set models to training mode generator.train() discriminator.train() for epoch in range(NUM_EPOCHS): for step, batch in enumerate(dataloader): # --- DISCRIMINATOR TRAINING STEP (D_STEPS times) --- for _ in range(D_STEPS): optimizer_d.zero_grad() # 1. Process REAL Data real_input_ids = batch['input_ids'] real_attention_mask = batch['attention_mask'] # Target: 1 (Real) real_labels = torch.ones(real_input_ids.size(0), 1).to(accelerator.device) # Get discriminator prediction for real data # We classify the full sequence (CLS token's output is used by RoBERTa's classification head) real_logits = discriminator(real_input_ids, attention_mask=real_attention_mask) real_loss = loss_fn(real_logits[:, 1].unsqueeze(-1), real_labels) # Use logit for label 1 (Real) # 2. Process FAKE (Generated) Data # Generate text using GPT-2. We use 'no_grad' since we don't want to calculate # gradients for the generator during the D step. with torch.no_grad(): # Generate text. 'max_length' ensures the generated text is the same size as real data. generated_ids = generator.generate( real_input_ids, max_length=SEQUENCE_LENGTH, do_sample=True, top_k=50, top_p=0.95, pad_token_id=tokenizer.eos_token_id ) # Get generated text and attention mask fake_input_ids = generated_ids # RoBERTa's tokenizer automatically handles attention masking based on the pad token ID fake_attention_mask = (fake_input_ids != tokenizer.pad_token_id).int() # Target: 0 (Fake) fake_labels = torch.zeros(fake_input_ids.size(0), 1).to(accelerator.device) # Get discriminator prediction for fake data fake_logits = discriminator(fake_input_ids, attention_mask=fake_attention_mask) # Use logit for label 1 (Real) but target is 0 (Fake). fake_loss = loss_fn(fake_logits[:, 1].unsqueeze(-1), fake_labels) # 3. Total Discriminator Loss and Update d_loss = real_loss + fake_loss # Backpropagate and update accelerator.backward(d_loss) optimizer_d.step() # --- GENERATOR TRAINING STEP (G_STEPS times) --- # Generator aims to make D classify its output as REAL (target 1) for _ in range(G_STEPS): optimizer_g.zero_grad() # Generate new fake data for the G step # We need gradients for this step, so no 'no_grad()' generated_ids = generator.generate( real_input_ids, max_length=SEQUENCE_LENGTH, do_sample=True, top_k=50, top_p=0.95, pad_token_id=tokenizer.eos_token_id ) fake_input_ids = generated_ids fake_attention_mask = (fake_input_ids != tokenizer.pad_token_id).int() # Target for Generator: 1 (it wants the discriminator to think the text is Real) generator_target_labels = torch.ones(fake_input_ids.size(0), 1).to(accelerator.device) # Get discriminator prediction for the generated text # We detach the discriminator's forward pass to prevent gradient updates to D during G step discriminator_logits = discriminator(fake_input_ids.detach(), attention_mask=fake_attention_mask.detach()) # Generator Loss: BCE loss where the target is 1 (Real) # The generator is being updated to minimize this loss, meaning its output # should drive the discriminator's output closer to 1. g_loss = loss_fn(discriminator_logits[:, 1].unsqueeze(-1), generator_target_labels) # Backpropagate and update accelerator.backward(g_loss) optimizer_g.step() # --- Logging and Reporting --- if (step + 1) % 50 == 0: # Calculate Discriminator Accuracy for monitoring # Predictions are based on which logit is higher (0 or 1) d_real_preds = (real_logits[:, 1] > real_logits[:, 0]).float().mean() d_fake_preds = (fake_logits[:, 1] < fake_logits[:, 0]).float().mean() d_accuracy = (d_real_preds + d_fake_preds) / 2 # G's success (how often D thinks the fake is real) g_success_rate = (discriminator_logits[:, 1] > discriminator_logits[:, 0]).float().mean() logger.info( f"Epoch {epoch+1}/{NUM_EPOCHS}, Step {step+1}/{len(dataloader)} | " f"D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f} | " f"D Acc: {d_accuracy.item():.2f} | G Success: {g_success_rate.item():.2f}" ) # --- End of Epoch --- logger.info(f"--- Epoch {epoch+1} finished. Generating sample text. ---") # Simple evaluation by generating text generator.eval() prompt = "Finetuning large language models in an adversarial setting is" input_ids = tokenizer.encode(prompt, return_tensors="pt").to(accelerator.device) sample_output = generator.generate( input_ids, max_length=50, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95, pad_token_id=tokenizer.eos_token_id ) decoded_output = tokenizer.decode(sample_output[0], skip_special_tokens=True) logger.info(f"Sample Output: {decoded_output}") generator.train() # Save the fine-tuned Generator model accelerator.wait_for_everyone() unwrapped_generator = accelerator.unwrap_model(generator) unwrapped_generator.save_pretrained("./finetuned_gpt2_gan_generator") logger.info("Fine-tuning complete. Generator saved to ./finetuned_gpt2_gan_generator") if __name__ == "__main__": # Note: To run this script, you typically need to use the 'accelerate launch' command: # accelerate launch your_script_name.py # Since this is a self-contained script in this environment, we call the function directly. train_gan()