Spaces:
Runtime error
Runtime error
| 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() |