Upload 2 files

distill.py

@@ -0,0 +1,264 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset, random_split
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+from typing import List, Optional
+import argparse
+import os
+import json
+import jsonlines
+from tqdm import tqdm
+from torch.cuda.amp import autocast, GradScaler
+from torch.utils.tensorboard import SummaryWriter
+
+# Set up device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class CustomDataset(Dataset):
+    def __init__(self, inputs, labels):
+        self.inputs = inputs
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, idx):
+        return {'input_ids': self.inputs[idx], 'labels': self.labels[idx]}
+
+def load_filtered_dataset(dataset_name: str, config: str, queries: Optional[List[str]] = None):
+    dataset = load_dataset(dataset_name, config)
+    if queries:
+        def filter_func(examples):
+            return any(query.lower() in examples["text"].lower() for query in queries)
+        dataset = dataset.filter(filter_func, batched=True)
+    return dataset
+
+def prepare_data(tokenizer, dataset, max_length, batch_size):
+    # Tokenize the inputs and labels
+    tokenized_inputs = tokenizer(dataset["train"]["text"], return_tensors="pt", padding=True, truncation=True, max_length=max_length)
+    tokenized_labels = tokenizer(dataset["train"]["text"], return_tensors="pt", padding=True, truncation=True, max_length=max_length)
+    
+    # Create custom dataset
+    custom_dataset = CustomDataset(tokenized_inputs["input_ids"], tokenized_labels["input_ids"])
+    
+    # Split into training and validation sets
+    train_size = int(0.9 * len(custom_dataset))
+    val_size = len(custom_dataset) - train_size
+    train_dataset, val_dataset = random_split(custom_dataset, [train_size, val_size])
+    
+    # Create DataLoaders
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
+    
+    return train_loader, val_loader
+
+def train_step(teacher, student, data_loader, optimizer, criterion, scaler, temperature=2.0):
+    teacher.eval()
+    student.train()
+    total_loss = 0
+
+    for batch in tqdm(data_loader, desc="Training"):
+        inputs = batch["input_ids"].to(device)
+        labels = batch["labels"].to(device)
+
+        with autocast():
+            with torch.no_grad():
+                teacher_outputs = teacher(inputs).logits
+                teacher_logits = teacher_outputs / temperature
+
+            student_outputs = student(inputs).logits
+            student_logits = student_outputs / temperature
+
+            # Compute KL Divergence Loss
+            loss = criterion(nn.functional.log_softmax(student_logits, dim=-1), nn.functional.softmax(teacher_logits, dim=-1))
+            loss = loss * (temperature ** 2)  # Scale loss by temperature squared
+
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad()
+
+        total_loss += loss.item()
+
+    avg_loss = total_loss / len(data_loader)
+    return avg_loss
+
+def validate(teacher, student, data_loader, criterion, temperature=2.0):
+    teacher.eval()
+    student.eval()
+    total_loss = 0
+
+    with torch.no_grad():
+        for batch in tqdm(data_loader, desc="Validation"):
+            inputs = batch["input_ids"].to(device)
+            labels = batch["labels"].to(device)
+
+            teacher_outputs = teacher(inputs).logits
+            teacher_logits = teacher_outputs / temperature
+
+            student_outputs = student(inputs).logits
+            student_logits = student_outputs / temperature
+
+            loss = criterion(nn.functional.log_softmax(student_logits, dim=-1), nn.functional.softmax(teacher_logits, dim=-1))
+            loss = loss * (temperature ** 2)
+
+            total_loss += loss.item()
+
+    avg_loss = total_loss / len(data_loader)
+    return avg_loss
+
+def save_checkpoint(state, save_dir, epoch):
+    os.makedirs(save_dir, exist_ok=True)
+    checkpoint_path = os.path.join(save_dir, f'checkpoint_epoch_{epoch}.pt')
+    torch.save(state, checkpoint_path)
+    print(f"Checkpoint saved at {checkpoint_path}")
+
+def load_checkpoint(model, optimizer, scheduler, scaler, save_dir, epoch):
+    checkpoint_path = os.path.join(save_dir, f'checkpoint_epoch_{epoch}.pt')
+    if os.path.isfile(checkpoint_path):
+        checkpoint = torch.load(checkpoint_path)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        print(f"Loaded checkpoint from {checkpoint_path}")
+    else:
+        print(f"No checkpoint found at {checkpoint_path}")
+
+def distill_model(
+    teacher_model_name: str,
+    student_model_name: str,
+    dataset_name: str,
+    config: str,
+    distill_full_model: bool = True,
+    query_terms: Optional[List[str]] = None,
+    num_epochs: int = 3,
+    batch_size: int = 4,
+    max_length: int = 128,
+    learning_rate: float = 5e-5,
+    temperature: float = 2.0,
+    save_path: str = "./distilled_model",
+    log_dir: str = "./logs",
+    checkpoint_dir: str = "./checkpoints",
+    early_stopping_patience: int = 3
+):
+    # Initialize TensorBoard writer
+    writer = SummaryWriter(log_dir=log_dir)
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load teacher and student models
+    teacher = AutoModelForCausalLM.from_pretrained(teacher_model_name).to(device)
+    student = AutoModelForCausalLM.from_pretrained(student_model_name).to(device)
+    
+    # Optionally freeze teacher model parameters
+    for param in teacher.parameters():
+        param.requires_grad = False
+    
+    # Load and prepare dataset
+    if distill_full_model:
+        dataset = load_dataset(dataset_name, config)
+    else:
+        dataset = load_filtered_dataset(dataset_name, config, query_terms)
+    
+    train_loader, val_loader = prepare_data(tokenizer, dataset, max_length, batch_size)
+    
+    # Define optimizer, scheduler, and scaler for mixed precision
+    optimizer = optim.AdamW(student.parameters(), lr=learning_rate)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
+    scaler = GradScaler()
+    
+    # Define loss criterion
+    criterion = nn.KLDivLoss(reduction="batchmean")
+    
+    best_val_loss = float('inf')
+    epochs_no_improve = 0
+    
+    # Training loop
+    for epoch in range(1, num_epochs + 1):
+        print(f"\nEpoch {epoch}/{num_epochs}")
+        print("-" * 20)
+        
+        # Training
+        train_loss = train_step(teacher, student, train_loader, optimizer, criterion, scaler, temperature)
+        print(f"Training Loss: {train_loss:.4f}")
+        writer.add_scalar("Loss/Train", train_loss, epoch)
+        
+        # Validation
+        val_loss = validate(teacher, student, val_loader, criterion, temperature)
+        print(f"Validation Loss: {val_loss:.4f}")
+        writer.add_scalar("Loss/Validation", val_loss, epoch)
+        
+        # Check for improvement
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            epochs_no_improve = 0
+            # Save the best model
+            save_checkpoint({
+                'epoch': epoch,
+                'model_state_dict': student.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'scheduler_state_dict': scheduler.state_dict(),
+                'scaler_state_dict': scaler.state_dict(),
+                'best_val_loss': best_val_loss
+            }, checkpoint_dir, epoch)
+            # Save the model as the best one
+            student.save_pretrained(save_path)
+            tokenizer.save_pretrained(save_path)
+            print(f"Best model saved at epoch {epoch}")
+        else:
+            epochs_no_improve += 1
+            print(f"No improvement in validation loss for {epochs_no_improve} epoch(s)")
+            if epochs_no_improve >= early_stopping_patience:
+                print("Early stopping triggered")
+                break
+        
+        # Step the scheduler
+        scheduler.step()
+    
+    writer.close()
+    print("\nDistillation completed.")
+
+def main():
+    parser = argparse.ArgumentParser(description="Distill a large LLM into a smaller one.")
+    parser.add_argument("--teacher_model_name", type=str, required=True, help="Name of the teacher model")
+    parser.add_argument("--student_model_name", type=str, required=True, help="Name of the student model")
+    parser.add_argument("--dataset_name", type=str, required=True, help="Name of the dataset")
+    parser.add_argument("--config", type=str, default=None, help="Dataset configuration (e.g., 'wikitext-2-raw-v1')")
+    parser.add_argument("--distill_full_model", action="store_true", help="Whether to distill the full model or not")
+    parser.add_argument("--query_terms", type=str, nargs="+", help="Query terms for filtering the dataset")
+    parser.add_argument("--num_epochs", type=int, default=3, help="Number of epochs")
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size")
+    parser.add_argument("--max_length", type=int, default=128, help="Maximum sequence length")
+    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
+    parser.add_argument("--temperature", type=float, default=2.0, help="Distillation temperature")
+    parser.add_argument("--save_path", type=str, default="./distilled_model", help="Path to save the distilled model")
+    parser.add_argument("--log_dir", type=str, default="./logs", help="Directory for TensorBoard logs")
+    parser.add_argument("--checkpoint_dir", type=str, default="./checkpoints", help="Directory to save checkpoints")
+    parser.add_argument("--early_stopping_patience", type=int, default=3, help="Early stopping patience")
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = main()
+    distill_model(
+        teacher_model_name=args.teacher_model_name,
+        student_model_name=args.student_model_name,
+        dataset_name=args.dataset_name,
+        config=args.config,
+        distill_full_model=args.distill_full_model,
+        query_terms=args.query_terms,
+        num_epochs=args.num_epochs,
+        batch_size=args.batch_size,
+        max_length=args.max_length,
+        learning_rate=args.learning_rate,
+        temperature=args.temperature,
+        save_path=args.save_path,
+        log_dir=args.log_dir,
+        checkpoint_dir=args.checkpoint_dir,
+        early_stopping_patience=args.early_stopping_patience
+    )

main_menu_new.py

@@ -0,0 +1,191 @@
+# main_menu.py
+
+import argparse
+import sys
+import os
+from train_agent import train_agent
+from test_agent import TestAgent, run_test_session
+from twisted.internet import reactor, task
+from lightbulb_custom import main as lightbulb_custom_main
+from distillation_pipeline import distill_model  # Import the distillation function
+from transformers import logging
+
+# Suppress transformers warnings for cleaner output
+logging.set_verbosity_error()
+
+def parse_main_args():
+    parser = argparse.ArgumentParser(description="Main Menu for Selecting Tasks")
+    
+    # Task selection
+    parser.add_argument('--task', type=str, choices=[
+                        'train_llm_world', 
+                        'train_agent', 
+                        'test_agent', 
+                        'inference_llm', 
+                        'inference_world_model', 
+                        'advanced_inference',
+                        'distill_full_model',       # New option for full model distillation
+                        'distill_domain_specific'   # New option for selective distillation
+                    ], 
+                        required=True, 
+                        help='Choose task to execute: train_llm_world, train_agent, test_agent, inference_llm, inference_world_model, advanced_inference, distill_full_model, distill_domain_specific')
+    
+    # Common arguments
+    parser.add_argument('--model_name', type=str, default='gpt2', help='Pretrained model name for LLM')
+    parser.add_argument('--student_model_name', type=str, default='distilgpt2', help='Name of the student model for distillation')
+    parser.add_argument('--dataset_name', type=str, default='wikitext', help='Dataset name for training')
+    parser.add_argument('--dataset_config', type=str, default='wikitext-2-raw-v1', help='Dataset configuration name')
+    parser.add_argument('--batch_size', type=int, default=4, help='Batch size for training')
+    parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs for training')
+    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length for training')
+    parser.add_argument('--temperature', type=float, default=2.0, help='Distillation temperature')
+    parser.add_argument('--learning_rate', type=float, default=5e-5, help='Learning rate')
+    
+    # Distillation-specific arguments
+    parser.add_argument('--save_path', type=str, default="./distilled_model", help="Path to save the distilled model")
+    parser.add_argument('--log_dir', type=str, default="./logs", help="Directory for TensorBoard logs")
+    parser.add_argument('--checkpoint_dir', type=str, default="./checkpoints", help="Directory to save checkpoints")
+    parser.add_argument('--early_stopping_patience', type=int, default=3, help="Early stopping patience")
+    
+    # Inference-specific arguments
+    parser.add_argument('--query', type=str, default='', help='Query for the test_agent or inference tasks')
+    parser.add_argument('--inference_mode', type=str, choices=['without_world_model', 'world_model', 'world_model_tree_of_thought'], help='Inference mode')
+    parser.add_argument('--beam_size', type=int, default=5, help='Beam size for beam search during inference')
+    parser.add_argument('--n_tokens_predict', type=int, default=3, help='Number of tokens to predict at each step during inference')
+    parser.add_argument('--mcts_iterations', type=int, default=10, help='Number of MCTS iterations during inference')
+    parser.add_argument('--mcts_exploration_constant', type=float, default=1.414, help='Exploration constant for MCTS during inference')
+    
+    # Distillation-specific arguments
+    parser.add_argument('--distill_full_model', action="store_true", help="Whether to distill the full model or not")
+    parser.add_argument('--query_terms', type=str, nargs="+", help="Query terms for domain-specific distillation")
+    
+    # Load model for inference
+    parser.add_argument('--load_model', type=str, help='Path to load the distilled model for inference')
+    
+    return parser.parse_args()
+
+def main():
+    # Parse arguments for the main function
+    args = parse_main_args()
+    
+    # Execute tasks based on user input
+    if args.task == 'train_llm_world':
+        print("Starting LLM and World Model Training...")
+        # Directly call the world model main function with appropriate arguments
+        sys.argv = [
+            'lightbulb_custom.py', 
+            '--mode', 'train', 
+            '--model_name', args.model_name, 
+            '--dataset_name', args.dataset_name, 
+            '--dataset_config', args.dataset_config,
+            '--batch_size', str(args.batch_size), 
+            '--num_epochs', str(args.num_epochs),
+            '--max_length', str(args.max_length)
+        ]
+        lightbulb_custom_main()
+    
+    elif args.task == 'train_agent':
+        print("Starting Agent Training...")
+        # Call the train_agent function from train_agent.py using Twisted reactor
+        d = task.deferLater(reactor, 0, train_agent)
+        d.addErrback(lambda failure: print(f"An error occurred: {failure}", exc_info=True))
+        d.addBoth(lambda _: reactor.stop())
+        reactor.run()
+    
+    elif args.task == 'test_agent':
+        print("Starting Test Agent...")
+        test_agent = TestAgent()
+        if args.query:
+            # Directly process a single query
+            result = test_agent.process_query(args.query)
+            print("\nAgent's response:")
+            print(result)
+        else:
+            # Run the interactive session
+            reactor.callWhenRunning(run_test_session)
+            reactor.run()
+    
+    elif args.task in ['inference_llm', 'inference_world_model', 'advanced_inference']:
+        print("Starting Inference Task...")
+        # Prepare the arguments for lightbulb_custom.py based on the selected inference task
+    
+        # Map the main_menu task to lightbulb_custom.py's inference_mode
+        inference_mode_map = {
+            'inference_llm': 'without_world_model',
+            'inference_world_model': 'world_model',
+            'advanced_inference': 'world_model_tree_of_thought'
+        }
+    
+        selected_inference_mode = inference_mode_map.get(args.task, 'world_model_tree_of_thought')
+    
+        # Construct sys.argv for lightbulb_custom.py
+        lightbulb_inf_args = [
+            'lightbulb_custom.py',
+            '--mode', 'inference',
+            '--model_name', args.model_name,
+            '--query', args.query,
+            '--max_length', str(args.max_length),
+            '--inference_mode', selected_inference_mode,
+            '--beam_size', str(args.beam_size),
+            '--n_tokens_predict', str(args.n_tokens_predict),
+            '--mcts_iterations', str(args.mcts_iterations),
+            '--mcts_exploration_constant', str(args.mcts_exploration_constant)
+        ]
+    
+        # Include additional arguments if they exist
+        if args.load_model:
+            lightbulb_inf_args += ['--load_model', args.load_model]
+    
+        # Update sys.argv and call the inference main function
+        sys.argv = lightbulb_inf_args
+        lightbulb_custom_main()
+    
+    elif args.task == 'distill_full_model':
+        print("Starting Full Model Distillation...")
+        distill_model(
+            teacher_model_name=args.model_name,
+            student_model_name=args.student_model_name,
+            dataset_name=args.dataset_name,
+            config=args.dataset_config,
+            distill_full_model=True,
+            query_terms=None,
+            num_epochs=args.num_epochs,
+            batch_size=args.batch_size,
+            max_length=args.max_length,
+            learning_rate=args.learning_rate,
+            temperature=args.temperature,
+            save_path=args.save_path,
+            log_dir=args.log_dir,
+            checkpoint_dir=args.checkpoint_dir,
+            early_stopping_patience=args.early_stopping_patience
+        )
+    
+    elif args.task == 'distill_domain_specific':
+        print("Starting Domain-Specific Distillation...")
+        if not args.query_terms:
+            print("Error: --query_terms must be provided for domain-specific distillation.")
+            sys.exit(1)
+        distill_model(
+            teacher_model_name=args.model_name,
+            student_model_name=args.student_model_name,
+            dataset_name=args.dataset_name,
+            config=args.dataset_config,
+            distill_full_model=False,
+            query_terms=args.query_terms,
+            num_epochs=args.num_epochs,
+            batch_size=args.batch_size,
+            max_length=args.max_length,
+            learning_rate=args.learning_rate,
+            temperature=args.temperature,
+            save_path=args.save_path,
+            log_dir=args.log_dir,
+            checkpoint_dir=args.checkpoint_dir,
+            early_stopping_patience=args.early_stopping_patience
+        )
+    
+    else:
+        print(f"Unknown task: {args.task}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()