app.py

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()