Create app.py

app.py

@@ -0,0 +1,261 @@
+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. Model and Tokenizer Initialization ---
+
+# Initialize the tokenizers and add padding token for GPT-2
+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
+
+# Initialize Models
+logger.info(f"Loading Generator ({GPT2_MODEL_NAME}) and Discriminator ({ROBERTA_MODEL_NAME})...")
+generator = GPT2LMHeadModel.from_pretrained(GPT2_MODEL_NAME)
+# Ensure the GPT-2 model's classification head (if used) is configured for the discriminator's vocabulary.
+# In a pure GAN, the generator just generates text, so we only need the LMHead.
+
+discriminator = Discriminator(ROBERTA_MODEL_NAME)
+
+# --- 3. Data Preprocessing ---
+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 Accelerator for mixed-precision and distributed training handling
+    accelerator = Accelerator()
+    
+    dataloader = load_and_prepare_data()
+
+    # Define optimizers
+    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()