Update distill.py

distill.py

@@ -1,264 +1,838 @@
-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
-    )
+import argparse
+import math
+import os
+import sys
+import json
+import jsonlines
+import copy
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset, random_split
+from torch.cuda.amp import autocast, GradScaler
+from torch.utils.tensorboard import SummaryWriter
+
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+from tqdm import tqdm
+
+# Set up device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ======================================
+# Import Custom Components from lightbulb_custom
+# ======================================
+from lightbulb_custom import (
+    RotaryPositionalEncoding,
+    MultiHeadAttention,
+    MoE,
+    TransformerBlock,
+    Transformer,
+    InfoNCE_Loss,
+    CovarianceRegularization,
+    DynamicsPerformanceLoss,
+    ThoughtConsistencyLoss,
+    PolicyValueJointLoss,
+    ActionDiversityReward,
+    ExpectedThoughtValueLoss,
+    ExplorationRegularization,
+    KL_DivergenceLoss,
+    ActionEncoder,
+    RepresentationNetwork,
+    DynamicsNetwork,
+    PredictionNetwork,
+    ThoughtNode,
+    MCTS,
+    State
+)
+
+# ==========================
+# Custom Dataset Definition
+# ==========================
+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]}
+
+# ================================
+# Utility Functions for Data Loading
+# ================================
+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 text.lower() for query in queries) for text in examples["text"]]
+        dataset = dataset.filter(filter_func, batched=True)
+    return dataset
+
+def load_custom_data_from_files(file_paths):
+    custom_data = []
+    for file_path in file_paths:
+        if file_path.endswith('.json'):
+            with open(file_path, 'r') as f:
+                data = json.load(f)
+                if isinstance(data, list):
+                    custom_data.extend(data)
+                else:
+                    custom_data.append(data)
+        elif file_path.endswith('.jsonl'):
+            with jsonlines.open(file_path) as reader:
+                custom_data.extend(reader)
+    return custom_data
+
+def preprocess_custom_data(data_list):
+    processed_data = []
+    for item in data_list:
+        # Check if the item is a string (JSON)
+        if isinstance(item, str):
+            try:
+                item = json.loads(item)
+            except json.JSONDecodeError:
+                print(f"Failed to parse JSON: {item[:100]}...")  # Print first 100 chars for debugging
+                continue  # Skip this item if it's not valid JSON
+
+        # Process query and content
+        query = item.get('query', '')
+        content = item.get('content', '')
+        if content == "RAG response generation failed.":
+            content = ""
+
+        # Combine query and content
+        combined_text = f"Query: {query} Content: {content}"
+
+        # Process numerical data (assuming these are available in the item dict)
+        episode_reward = item.get('episode_reward', 0)
+        loss = item.get('loss', 0)
+        cosine_similarity = item.get('cosine_similarity', 0)
+        rag_performance = item.get('rag_performance', 0)
+        ranking_model_performance = item.get('ranking_model_performance', 0)
+
+        # Create a dictionary with processed data
+        processed_item = {
+            'text': combined_text,
+            'episode_reward': episode_reward,
+            'loss': loss,
+            'cosine_similarity': cosine_similarity,
+            'rag_performance': rag_performance,
+            'ranking_model_performance': ranking_model_performance
+        }
+
+        processed_data.append(processed_item)
+
+    return processed_data
+
+def load_custom_data(args, tokenizer, custom_data):
+    # Preprocess the custom data
+    processed_data = preprocess_custom_data(custom_data)
+
+    # Create a custom dataset
+    class CustomDatasetProcessed(torch.utils.data.Dataset):
+        def __init__(self, data, tokenizer, max_length):
+            self.data = data
+            self.tokenizer = tokenizer
+            self.max_length = max_length
+
+        def __len__(self):
+            return len(self.data)
+
+        def __getitem__(self, idx):
+            item = self.data[idx]
+            encoded = self.tokenizer.encode_plus(
+                item['text'],
+                max_length=self.max_length,
+                padding='max_length',
+                truncation=True,
+                return_tensors='pt'
+            )
+            return {
+                'input_ids': encoded['input_ids'].squeeze(),
+                'attention_mask': encoded['attention_mask'].squeeze(),
+                'episode_reward': torch.tensor(item['episode_reward'], dtype=torch.float),
+                'loss': torch.tensor(item['loss'], dtype=torch.float),
+                'cosine_similarity': torch.tensor(item['cosine_similarity'], dtype=torch.float),
+                'rag_performance': torch.tensor(item['rag_performance'], dtype=torch.float),
+                'ranking_model_performance': torch.tensor(item['ranking_model_performance'], dtype=torch.float)
+            }
+
+    # Create dataset and dataloader
+    dataset = CustomDatasetProcessed(processed_data, tokenizer, args.max_length)
+
+    # Split the dataset into train and eval
+    train_size = int(0.8 * len(dataset))
+    eval_size = len(dataset) - train_size
+    train_dataset, eval_dataset = random_split(dataset, [train_size, eval_size])
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=4
+    )
+    eval_loader = DataLoader(
+        eval_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=4
+    )
+
+    return train_loader, eval_loader
+
+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,
+        shuffle=True,
+        batch_size=batch_size,
+        num_workers=4,
+        pin_memory=True
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        shuffle=False,
+        batch_size=batch_size,
+        num_workers=4,
+        pin_memory=True
+    )
+
+    return train_loader, val_loader
+
+# ==========================
+# Training and Validation Functions
+# ==========================
+
+def save_all_models(transformer_model, representation_network, dynamics_network, prediction_network, action_encoder, save_dir, epoch):
+    """
+    Save all models to the specified directory.
+    Args:
+        transformer_model (nn.Module): Transformer model.
+        representation_network (nn.Module): Representation network.
+        dynamics_network (nn.Module): Dynamics network.
+        prediction_network (nn.Module): Prediction network.
+        action_encoder (nn.Module): Action encoder.
+        save_dir (str): Directory to save the models.
+        epoch (int): Current epoch number.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+
+    torch.save(transformer_model.state_dict(), os.path.join(save_dir, f'transformer_model_epoch_{epoch}.pt'))
+    torch.save(representation_network.state_dict(), os.path.join(save_dir, f'representation_network_epoch_{epoch}.pt'))
+    torch.save(dynamics_network.state_dict(), os.path.join(save_dir, f'dynamics_network_epoch_{epoch}.pt'))
+    torch.save(prediction_network.state_dict(), os.path.join(save_dir, f'prediction_network_epoch_{epoch}.pt'))
+    torch.save(action_encoder.state_dict(), os.path.join(save_dir, f'action_encoder_epoch_{epoch}.pt'))
+
+    print(f"All models saved for epoch {epoch}.")
+
+def train_epoch_world_model(world_model_components, train_loader, optimizer, scheduler, scaler, args, model_transformer, state_dim, embed_dim, input_dim):
+    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent, model_transformer = world_model_components
+    representation_network.train()
+    dynamics_network.train()
+    prediction_network.train()
+    action_encoder.train()
+    ppo_agent.policy_network.train()
+
+    total_loss = 0.0
+    optimizer.zero_grad()
+    print(f"Starting World Model training epoch with {len(train_loader)} batches...")
+
+    for i, batch in enumerate(train_loader):
+        print(f"Processing batch {i+1}/{len(train_loader)}...")
+
+        # Move batches to the device
+        src_batch = batch['input_ids'].to(device)
+        tgt_batch = batch['labels'].to(device)
+
+        with torch.cuda.amp.autocast():
+            print("Forward pass through Transformer (frozen)...")
+            with torch.no_grad():
+                transformer_output = model_transformer(src_batch, tgt_batch[:, :-1])
+
+            # World Model - Representation
+            state_representation = representation_network(transformer_output)
+
+            # For simplicity, let's assume true actions are provided (e.g., next tokens)
+            true_actions = tgt_batch[:, :-1]
+            print(f"True actions shape: {true_actions.shape}")
+            action_sequences = true_actions
+
+            # Get action embeddings
+            action_embeddings = action_encoder(action_sequences)
+            print(f"Action embeddings shape: {action_embeddings.shape}")
+
+            # Apply dynamics network
+            predicted_next_state_batch = dynamics_network(state_representation, action_embeddings)
+            print(f"Predicted next state batch shape: {predicted_next_state_batch.shape}")
+
+            # Prediction Network - Policy logits and value
+            policy_logits, value_estimates = prediction_network(predicted_next_state_batch)
+
+            # Define true_policy and true_value as placeholders on the GPU
+            true_policy = F.one_hot(true_actions, num_classes=input_dim).float()
+            true_value = torch.zeros_like(value_estimates).to(device)
+
+            # Compute individual losses
+            ppo_loss = ppo_agent.compute_loss(
+                state_representation,
+                torch.zeros_like(true_actions, dtype=torch.float32).to(device),
+                true_actions,
+                torch.zeros_like(value_estimates, dtype=torch.float32).to(device),
+                torch.zeros_like(value_estimates, dtype=torch.float32).to(device)
+            )
+
+            info_nce = InfoNCE_Loss()(state_representation.reshape(-1, state_dim),
+                                      F.dropout(state_representation.reshape(-1, state_dim), p=0.1, training=True))
+
+            covariance = CovarianceRegularization()(predicted_next_state_batch.view(-1, predicted_next_state_batch.size(-1)))
+            dynamics_loss = DynamicsPerformanceLoss()(state_representation, predicted_next_state_batch)
+
+            perturbed_next_state = predicted_next_state_batch + torch.randn_like(predicted_next_state_batch) * 0.01
+            thought_loss = ThoughtConsistencyLoss()(predicted_next_state_batch, perturbed_next_state)
+
+            pv_loss = PolicyValueJointLoss()(policy_logits, true_policy, value_estimates.squeeze(-1), true_value.squeeze(-1))
+            action_diversity = ActionDiversityReward()(action_embeddings.view(-1, embed_dim))
+
+            mcts_best_values = torch.zeros(true_actions.size(0)).to(device)
+            etv = ExpectedThoughtValueLoss()(mcts_best_values)
+
+            visit_counts = torch.ones(true_actions.size(0), policy_logits.size(-1)).to(device)
+            exploration = ExplorationRegularization()(visit_counts)
+
+            old_policy = F.softmax(policy_logits.detach(), dim=-1)
+            new_policy = F.softmax(policy_logits, dim=-1)
+            kl_loss = KL_DivergenceLoss()(old_policy, new_policy)
+
+            # Total Loss
+            loss = (
+                ppo_loss +
+                info_nce +
+                covariance +
+                dynamics_loss +
+                thought_loss +
+                pv_loss +
+                action_diversity +
+                etv +
+                exploration +
+                kl_loss
+            )
+            loss = loss / args.accumulation_steps
+
+        print("Backward pass...")
+        scaler.scale(loss).backward()
+
+        if (i + 1) % args.accumulation_steps == 0 or (i + 1) == len(train_loader):
+            print("Gradient clipping...")
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(
+                [param for group in optimizer.param_groups for param in group['params']],
+                args.max_grad_norm
+            )
+
+            print("Optimizer step...")
+            scaler.step(optimizer)
+            scaler.update()
+
+            print("Zeroing gradients...")
+            optimizer.zero_grad()
+
+            print("Updating learning rate...")
+            scheduler.step()
+
+        total_loss += loss.item() * args.accumulation_steps
+
+        # Print individual losses and total loss for this batch
+        print(f"Batch {i+1} completed. Losses:")
+        print(f"  PPO Loss: {ppo_loss.item():.4f}")
+        print(f"  InfoNCE Loss: {info_nce.item():.4f}")
+        print(f"  Covariance Loss: {covariance.item():.4f}")
+        print(f"  Dynamics Loss: {dynamics_loss.item():.4f}")
+        print(f"  Thought Consistency Loss: {thought_loss.item():.4f}")
+        print(f"  Policy-Value Loss: {pv_loss.item():.4f}")
+        print(f"  Action Diversity Loss: {action_diversity.item():.4f}")
+        print(f"  Expected Thought Value Loss: {etv.item():.4f}")
+        print(f"  Exploration Loss: {exploration.item():.4f}")
+        print(f"  KL Divergence Loss: {kl_loss.item():.4f}")
+        print(f"  Total Loss: {loss.item():.4f}")
+
+    avg_loss = total_loss / len(train_loader)
+    print(f"World Model training epoch completed. Average loss: {avg_loss:.4f}")
+    return avg_loss
+
+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}")
+
+# ==========================
+# Inference Functions
+# ==========================
+
+def infer(query, world_model_components, root_thought_node, tokenizer, max_length=2000, inference_mode='world_model', beam_size=5, n_tokens_predict=3, mcts_iterations=10, exploration_constant=1.414):
+    """
+    Perform inference given a query, utilizing the Tree of Thought and MCTS with multi-token beam search.
+    Args:
+        query (str): The input query or prompt.
+        world_model_components (tuple): Tuple containing the model components.
+        root_thought_node (ThoughtNode): The root node of the Tree of Thought.
+        tokenizer (transformers.PreTrainedTokenizer): The tokenizer used.
+        max_length (int): Maximum length for the generated sequence.
+        inference_mode (str): Inference mode ('world_model', 'without_world_model', 'world_model_tree_of_thought')
+        beam_size (int): Size of the beam for beam search
+        n_tokens_predict (int): Number of tokens to predict at each step
+        mcts_iterations (int): Number of MCTS iterations
+        exploration_constant (float): Exploration constant for MCTS
+    Returns:
+        List[str] or str: The sequence of actions (thoughts) selected or generated text.
+    """
+    if inference_mode != 'world_model':
+        print("Inference mode other than 'world_model' not implemented yet.")
+        return ""
+
+    representation_network, dynamics_network, prediction_network, action_encoder, ppo_agent, model_transformer = world_model_components
+
+    # Tokenize and encode the query
+    input_ids = tokenizer.encode(query, return_tensors='pt').to(device)
+    attention_mask = (input_ids != tokenizer.pad_token_id).long()
+
+    # Use the world model components
+    with torch.no_grad():
+        transformer_output = model_transformer(input_ids, input_ids)
+        # Get the initial state representation
+        initial_representation = representation_network(transformer_output)  # Shape: (batch_size=1, seq_len, state_dim)
+        initial_representation = initial_representation[:, -1, :].unsqueeze(1)  # Shape: (batch_size=1, 1, state_dim)
+        initial_state = State(
+            representation=initial_representation,
+            dynamics_network=dynamics_network,
+            action_encoder=action_encoder,
+            thought_node=root_thought_node
+        )
+        # Use MCTS with Tree of Thought and multi-token beam search
+        mcts = MCTS(prediction_network, dynamics_network, action_encoder, num_iterations=mcts_iterations, exploration_constant=exploration_constant)
+
+        current_state = initial_state
+        thought_sequence = []
+
+        for _ in range(max_length // n_tokens_predict):
+            best_actions = mcts.search_with_beam(current_state)
+
+            thought_sequence.extend(best_actions)
+
+            # Apply the best actions to get the next state
+            for action in best_actions:
+                current_state = current_state.apply_action(action)
+
+            # Check if we've reached a leaf node (no further actions)
+            if len(current_state.thought_node.children) == 0:
+                break
+
+        return thought_sequence
+
+# ==========================
+# Main Training Function
+# ==========================
+
+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,
+    accumulation_steps: int = 1,
+    max_grad_norm: float = 1.0,
+    weight_decay: float = 0.01
+):
+    # Initialize TensorBoard writer
+    writer = SummaryWriter(log_dir=log_dir)
+
+    # Load tokenizer
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(teacher_model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    print("Tokenizer loaded successfully.")
+
+    # Load teacher model
+    print("Loading teacher model...")
+    teacher = AutoModelForCausalLM.from_pretrained(teacher_model_name).to(device)
+    print("Teacher model loaded successfully.")
+
+    if distill_full_model:
+        # Full World Model Distillation
+        print(f"Starting Full World Model Distillation into '{student_model_name}'.")
+
+        # Load or instantiate student model
+        print(f"Attempting to load student model '{student_model_name}'...")
+        try:
+            student = AutoModelForCausalLM.from_pretrained(student_model_name).to(device)
+            print(f"Student model '{student_model_name}' loaded successfully.")
+        except (OSError, ValueError) as e:
+            print(f"Student model '{student_model_name}' not found. Instantiating a new student model.")
+            # Instantiate a smaller pre-trained model as the student, e.g., distilgpt2
+            try:
+                student = AutoModelForCausalLM.from_pretrained('distilgpt2').to(device)
+                # Save the instantiated student model with the desired name
+                student.save_pretrained(save_path)
+                tokenizer.save_pretrained(save_path)
+                print(f"New student model '{student_model_name}' instantiated and saved to '{save_path}'.")
+            except Exception as inst_e:
+                print(f"Failed to instantiate and save student model: {inst_e}")
+                sys.exit(1)
+
+        # Optionally freeze teacher model parameters
+        for param in teacher.parameters():
+            param.requires_grad = False
+
+        # Load and prepare dataset
+        print(f"Loading full dataset '{dataset_name}' with config '{config}'...")
+        dataset = load_dataset(dataset_name, config)
+        train_loader, val_loader = prepare_data(tokenizer, dataset, max_length, batch_size)
+        print("Data loaded and preprocessed successfully.")
+
+        # Define optimizer, scheduler, and scaler for mixed precision
+        optimizer = optim.AdamW(student.parameters(), lr=learning_rate, weight_decay=weight_decay)
+        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("\nFull World Model Distillation completed.")
+
+    else:
+        # Standard Language Model Distillation
+        print(f"Starting Standard Language Model Distillation into '{student_model_name}'.")
+
+        if not query_terms:
+            print("Error: --query_terms must be provided for standard language model distillation.")
+            sys.exit(1)
+
+        # Load or instantiate student model
+        print(f"Attempting to load student model '{student_model_name}'...")
+        try:
+            student = AutoModelForCausalLM.from_pretrained(student_model_name).to(device)
+            print(f"Student model '{student_model_name}' loaded successfully.")
+        except (OSError, ValueError) as e:
+            print(f"Student model '{student_model_name}' not found. Instantiating a new student model.")
+            # Instantiate a smaller pre-trained model as the student, e.g., distilgpt2
+            try:
+                student = AutoModelForCausalLM.from_pretrained('distilgpt2').to(device)
+                # Save the instantiated student model with the desired name
+                student.save_pretrained(save_path)
+                tokenizer.save_pretrained(save_path)
+                print(f"New student model '{student_model_name}' instantiated and saved to '{save_path}'.")
+            except Exception as inst_e:
+                print(f"Failed to instantiate and save student model: {inst_e}")
+                sys.exit(1)
+
+        # Optionally freeze teacher model parameters
+        for param in teacher.parameters():
+            param.requires_grad = False
+
+        # Load and prepare custom dataset
+        print(f"Loading custom data files: {query_terms}")
+        custom_data = load_custom_data_from_files(query_terms)
+        train_loader, val_loader = load_custom_data(
+            args=argparse.Namespace(max_length=max_length),
+            tokenizer=tokenizer,
+            custom_data=custom_data
+        )
+        print("Custom data loaded and preprocessed successfully.")
+
+        # Define optimizer, scheduler, and scaler for mixed precision
+        optimizer = optim.AdamW(student.parameters(), lr=learning_rate, weight_decay=weight_decay)
+        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("\nStandard Language Model Distillation completed.")
+
+# ==========================
+# Argument Parsing
+# ==========================
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Distill a large LLM into a smaller one or a full language world model.")
+
+    # Required arguments
+    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")
+
+    # Dataset arguments
+    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')")
+
+    # Mode selection
+    parser.add_argument("--distill_full_model", action="store_true", help="Whether to distill into the full language world model")
+
+    # For standard distillation
+    parser.add_argument("--query_terms", type=str, nargs="+", help="Paths to custom data files for standard language model distillation")
+
+    # Training hyperparameters
+    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")
+
+    # Saving and logging
+    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")
+
+    # Early stopping
+    parser.add_argument("--early_stopping_patience", type=int, default=3, help="Early stopping patience")
+
+    # Gradient accumulation and optimization
+    parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
+    parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Maximum gradient norm for clipping")
+    parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay for optimizer")
+
+    return parser.parse_args()
+
+# ==========================
+# Main Function
+# ==========================
+
+def main():
+    args = parse_args()
+    print("Arguments parsed successfully.")
+
+    # Create save directories
+    os.makedirs(args.save_path, exist_ok=True)
+    os.makedirs(args.log_dir, exist_ok=True)
+    os.makedirs(args.checkpoint_dir, exist_ok=True)
+    print(f"Save directory created: {args.save_path}")
+    print(f"Log directory created: {args.log_dir}")
+    print(f"Checkpoint directory created: {args.checkpoint_dir}")
+
+    # Handle dataset loading based on distillation mode
+    if args.distill_full_model:
+        # Full World Model Distillation
+        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,  # Not used in this mode
+            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,
+            accumulation_steps=args.accumulation_steps,
+            max_grad_norm=args.max_grad_norm,
+            weight_decay=args.weight_decay
+        )
+    else:
+        # Standard Language Model Distillation
+        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,
+            accumulation_steps=args.accumulation_steps,
+            max_grad_norm=args.max_grad_norm,
+            weight_decay=args.weight_decay
+        )
+
+
+
+if __name__ == "__main__":
+    main()
+
+