train_withmtp.py

# ========================
#  Train with NTP + MTP
#  Updated for ChemQ3MTP structure
#  by gbyuvd
# ========================

# train_withmtp.py
import sys
import os
# Add the current directory to Python path so it can find your modules
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
import torch.nn as nn
import torch.nn.functional as F
import json
from typing import List, Union, Optional, Tuple, Dict, Any
from transformers.tokenization_utils_base import BatchEncoding
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset, DatasetDict, Dataset
import pandas as pd
from torch.utils.data import Dataset as TorchDataset, DataLoader, random_split
from sklearn.model_selection import train_test_split
from ranger21 import Ranger21
from tqdm.notebook import tqdm
from FastChemTokenizerHF import FastChemTokenizerSelfies
from ChemQ3MTP import ChemQ3MTPConfig, ChemQ3MTPForCausalLM  # This should now work
os.environ["TOKENIZERS_PARALLELISM"] = "false"
from transformers import TrainerCallback
import datetime

# Clear cache functions
def clear_cache():
    """Clear PyTorch and CUDA caches"""
    print("Clearing PyTorch and CUDA caches...")
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
        print("CUDA cache cleared")
    torch.backends.cudnn.benchmark = True  # Enable cuDNN optimization
    print("PyTorch cache cleared")

def clear_datasets_cache():
    """Clear datasets cache directory"""
    import shutil
    from datasets import disable_caching, enable_caching, get_cache_directory
    try:
        cache_dir = get_cache_directory()
        print(f"Clearing datasets cache at: {cache_dir}")
        if os.path.exists(cache_dir):
            shutil.rmtree(cache_dir)
            print("Datasets cache cleared")
    except:
        print("Could not clear datasets cache (may not exist)")

# ==============================
# Clear caches before starting
# ==============================
clear_cache()
# clear_datasets_cache()

# ==============================
# Load external configuration
# ==============================
with open("config.json", "r") as f:
    CONFIG = json.load(f)

TRAINING_CFG = CONFIG["training"]
MODEL_CFG = {k: v for k, v in CONFIG.items() 
             if k not in ["training", "generation", "model_type", "architectures"]}
GENERATION_CFG = CONFIG.get("generation", {})

# Training params
BATCH_SIZE = TRAINING_CFG["batch_size"]
NUM_EPOCHS = TRAINING_CFG["num_epochs"]
LEARNING_RATE = TRAINING_CFG["learning_rate"]
WEIGHT_DECAY = TRAINING_CFG["weight_decay"]
GRAD_ACCUM_STEPS = TRAINING_CFG["gradient_accumulation_steps"]
TOKENIZE_BATCH_SIZE = TRAINING_CFG["tokenize_batch_size"]
TRAIN_SPLIT_RATIO = TRAINING_CFG["train_split_ratio"]
VAL_SPLIT_RATIO = TRAINING_CFG["val_split_ratio"]
TEST_SPLIT_RATIO = TRAINING_CFG["test_split_ratio"]
INCLUDE_FOR_METRICS = TRAINING_CFG.get("include_for_metrics", ["input_ids", "attention_mask", "labels"])
# ==============================

class LossLoggerCallback(TrainerCallback):
    def __init__(self, log_file="training_losses.txt", with_timestamp=False):
        self.log_file = log_file
        self.with_timestamp = with_timestamp
        with open(self.log_file, "w") as f:
            if self.with_timestamp:
                f.write("time\tstep\tloss\teval_loss\n")
            else:
                f.write("step\tloss\teval_loss\n")

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None:
            return
        step = state.global_step
        loss = logs.get("loss")
        eval_loss = logs.get("eval_loss")

        with open(self.log_file, "a") as f:
            if self.with_timestamp:
                ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                f.write(f"{ts}\t{step}\t{loss if loss is not None else ''}\t{eval_loss if eval_loss is not None else ''}\n")
            else:
                f.write(f"{step}\t{loss if loss is not None else ''}\t{eval_loss if eval_loss is not None else ''}\n")


class CheckpointEvery10PercentCallback(TrainerCallback):
    """
    Custom callback to save checkpoints at 10% intervals of total training progress
    """
    def __init__(self, save_dir, total_steps):
        self.save_dir = save_dir
        self.total_steps = total_steps
        self.checkpoint_intervals = []
        # Calculate steps for 10% intervals (10%, 20%, 30%, ..., 100%)
        for i in range(1, 11):
            checkpoint_step = int(total_steps * i * 0.1)
            self.checkpoint_intervals.append(checkpoint_step)
        self.saved_checkpoints = set()
        print(f"Checkpoint intervals: {self.checkpoint_intervals}")

    def on_step_end(self, args, state, control, **kwargs):
        current_step = state.global_step
        
        # Check if we've reached a 10% checkpoint
        for checkpoint_step in self.checkpoint_intervals:
            if current_step == checkpoint_step and checkpoint_step not in self.saved_checkpoints:
                checkpoint_dir = f"{self.save_dir}/checkpoint_10percent_{current_step}"
                print(f"Saving 10% progress checkpoint at step {current_step} to {checkpoint_dir}")
                
                # Save model and tokenizer
                model = kwargs.get('model')
                tokenizer = kwargs.get('processing_class')  # or kwargs.get('tokenizer')
                
                if model is not None:
                    model.save_pretrained(checkpoint_dir)
                if tokenizer is not None:
                    tokenizer.save_pretrained(checkpoint_dir)
                
                # Also save training state
                if hasattr(kwargs.get('trainer'), 'save_state'):
                    kwargs['trainer'].save_state()
                
                self.saved_checkpoints.add(checkpoint_step)
                print(f"Checkpoint saved at step {current_step} ({current_step/self.total_steps*100:.1f}% completion)")
                break  # Only save one checkpoint per step


def tokenize_function(examples, tokenizer, max_length):
    """Tokenize function defined outside main to avoid closure issues"""
    batch_results = {"input_ids": [], "attention_mask": [], "labels": []}
    smiles_list = examples['SELFIES'] if isinstance(examples['SELFIES'], list) else [examples['SELFIES']]
    for smiles in smiles_list:
        tokenized = tokenizer(
            smiles,
            truncation=True,
            padding=False,
            max_length=max_length,
            return_tensors=None,
            add_special_tokens=True
        )
        input_ids = tokenized["input_ids"]
        attention_mask = tokenized["attention_mask"]
        labels = input_ids.copy()
        batch_results["input_ids"].append(input_ids)
        batch_results["attention_mask"].append(attention_mask)
        batch_results["labels"].append(labels)
    return batch_results


def main():
    # Clear cache at the beginning of main function too
    clear_cache()
    
    # --- Load the tokenizer ---
    tokenizer = FastChemTokenizerSelfies.from_pretrained("../selftok_core")

    out = tokenizer("[C] [=C] [Branch1]", return_tensors="pt")
    print(out.input_ids)
    print(out.attention_mask)
    out = out.to("cuda" if torch.cuda.is_available() else "cpu")
    print(out.input_ids.device)

    # --- Define config ---
    config = ChemQ3MTPConfig(  # Updated to use ChemQ3MTPConfig
        vocab_size=len(tokenizer),
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        **MODEL_CFG
    )

    model = ChemQ3MTPForCausalLM(config)  # Updated to use the new model class

    def count_parameters(model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)

    print(f"Enhanced model has {count_parameters(model):,} trainable parameters.")

    batch_size, seq_len = 2, 32
    dummy_input = torch.randint(
        low=0,
        high=len(tokenizer),
        size=(batch_size, seq_len),
        dtype=torch.long,
    )
    with torch.no_grad():
        outputs = model(dummy_input)
        logits = outputs.logits
    print(f"Input shape: {dummy_input.shape}")
    print(f"Logits shape: {logits.shape}")

    print("Loading dataset...")
    # Load dataset without streaming
    dataset = load_dataset(
        'csv',
        data_files='../data/chunk_1.csv',
        split='train'
    )

    print(f"Dataset loaded with {len(dataset)} samples")
    
    # Verify the correct file is loaded by checking first few samples
    print("First few samples from dataset:")
    for i in range(min(3, len(dataset))):
        sample = dataset[i]
        print(f"Sample {i}: {sample}")
        if 'SELFIES' in sample:
            print(f"First SELFIES: {sample['SELFIES']}")
        break

    print("Shuffling and splitting dataset...")
    # Shuffle the entire dataset first
    dataset = dataset.shuffle(seed=42)

    # Calculate split sizes
    total_lines = len(dataset)
    test_size = int(TEST_SPLIT_RATIO * total_lines)
    val_size = int(VAL_SPLIT_RATIO * total_lines)
    train_size = total_lines - test_size - val_size

    print(f"Total samples: {total_lines}")
    print(f"Split sizes - train: {train_size}, val: {val_size}, test: {test_size}")

    # Create splits using select
    train_dataset = dataset.select(range(0, train_size))
    val_dataset = dataset.select(range(train_size, train_size + val_size))
    test_dataset = dataset.select(range(train_size + val_size, total_lines))

    print(f"Dataset split: train={len(train_dataset)}, val={len(val_dataset)}, test={len(test_dataset)}")

    # Tokenize datasets using batched mapping with explicit parameters
    print("Tokenizing datasets...")
    
    # Define tokenize function with all parameters passed explicitly
    def tokenize_train(examples):
        return tokenize_function(examples, tokenizer, MODEL_CFG["max_position_embeddings"])
    
    def tokenize_val(examples):
        return tokenize_function(examples, tokenizer, MODEL_CFG["max_position_embeddings"])

    train_dataset = train_dataset.map(
        tokenize_train, 
        batched=True, 
        batch_size=TOKENIZE_BATCH_SIZE, 
        remove_columns=["SELFIES"],
        desc="Tokenizing train dataset"
    )
    val_dataset = val_dataset.map(
        tokenize_val, 
        batched=True, 
        batch_size=TOKENIZE_BATCH_SIZE, 
        remove_columns=["SELFIES"],
        desc="Tokenizing val dataset"
    )

    class EnhancedDataCollator:
        def __init__(self, tokenizer, pad_to_multiple_of=8):
            self.tokenizer = tokenizer
            self.pad_to_multiple_of = pad_to_multiple_of
        def __call__(self, features):
            max_length = max(len(f["input_ids"]) for f in features)
            if self.pad_to_multiple_of:
                max_length = ((max_length + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of) * self.pad_to_multiple_of
            batch = {"input_ids": [], "attention_mask": [], "labels": []}
            for feature in features:
                input_ids = feature["input_ids"]
                attention_mask = feature["attention_mask"]
                labels = feature["labels"]
                padding_length = max_length - len(input_ids)
                padded_input_ids = input_ids + [self.tokenizer.pad_token_id] * padding_length
                padded_attention_mask = attention_mask + [0] * padding_length
                padded_labels = labels + [-100] * padding_length
                batch["input_ids"].append(padded_input_ids)
                batch["attention_mask"].append(padded_attention_mask)
                batch["labels"].append(padded_labels)
            batch = {key: torch.tensor(values, dtype=torch.long) for key, values in batch.items()}
            return batch

    data_collator = EnhancedDataCollator(tokenizer, pad_to_multiple_of=8)

    def create_enhanced_optimizer(model_params):
        num_batches_per_epoch = len(train_dataset) // BATCH_SIZE
        optimizer_params = {
            'lr': LEARNING_RATE,
            'weight_decay': WEIGHT_DECAY,
            'use_adabelief': True,
            'use_cheb': False,
            'use_warmup': True,
            'use_madgrad': True,
            'num_epochs': NUM_EPOCHS,
            'using_gc': True,
            'warmdown_active': True,
            'num_batches_per_epoch': num_batches_per_epoch
        }
        return Ranger21(model_params, **optimizer_params)

    from torch.optim.lr_scheduler import LambdaLR
    class EnhancedCustomTrainer(Trainer):
        def create_optimizer(self):
            self.optimizer = create_enhanced_optimizer(self.model.parameters())
            return self.optimizer
        def create_scheduler(self, num_training_steps, optimizer=None):
            if optimizer is None:
                optimizer = self.optimizer
            self.lr_scheduler = LambdaLR(optimizer, lr_lambda=lambda step: 1.0)
            return self.lr_scheduler
        def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
            outputs = model(**inputs)
            loss = outputs.loss
            return (loss, outputs) if return_outputs else loss

    steps_per_epoch = len(train_dataset) // BATCH_SIZE
    total_steps = steps_per_epoch * NUM_EPOCHS

    training_args = TrainingArguments(
        output_dir='./chemq3minipret',
        max_steps=total_steps,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRAD_ACCUM_STEPS,
        logging_dir='./gptlo-1',
        logging_strategy="steps",
        logging_steps=max(1, steps_per_epoch // 4),
        eval_strategy="steps",
        eval_steps=max(1, steps_per_epoch // 4),
        save_strategy="steps",
        save_steps=steps_per_epoch,  # Save every epoch
        save_total_limit=1,
        dataloader_num_workers=0,
        dataloader_pin_memory=False,
        remove_unused_columns=False,
        prediction_loss_only=False,
        fp16=torch.cuda.is_available(),
        gradient_checkpointing=True,
        dataloader_drop_last=True,
        report_to=None,
        include_for_metrics=INCLUDE_FOR_METRICS,
    )

    print("Initializing enhanced trainer with MTP capabilities...")
    trainer = EnhancedCustomTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=data_collator,
        processing_class=tokenizer,
        callbacks=[
            LossLoggerCallback("training_losses.txt", with_timestamp=True),
            CheckpointEvery10PercentCallback("./chemq3minipret", total_steps)
        ]
    )

    model.set_mtp_training(True)
    print(" MTP training mode enabled")

    print("Starting enhanced training with MTP and Horizon Loss...")
    try:
        print("\n Phase 1: Warmup with standard Causal LM...")
        model.set_mtp_training(False)
        warmup_steps = max(1, total_steps // 5)
        
        # Update trainer args for warmup phase
        trainer.args.max_steps = warmup_steps
        trainer.train()
        print(f"\n Phase 1 completed. Warmup with {warmup_steps} steps finished.")
        
        print(f"\n Phase 2: Full MTP + Horizon Loss training...")
        print(f"Total training steps: {total_steps}")
        print(f"Training will save checkpoints at 10% intervals:")
        for i in range(1, 11):
            checkpoint_step = int(total_steps * i * 0.1)
            print(f"  - {i*10}%: Step {checkpoint_step}")
        
        model.set_mtp_training(True)
        # Reset max steps to total for the full training phase
        trainer.args.max_steps = total_steps
        trainer.train(resume_from_checkpoint=True)
        print("Enhanced training completed successfully!")
        trainer.save_model("./enhanced-qwen3-final")
        tokenizer.save_pretrained("./enhanced-qwen3-final")
        training_config = {
            "model_type": "ChemQ3MTPForCausalLM",
            "num_future_tokens": 3,
            "horizon_loss_enabled": True,
            "mtp_head_enabled": True,
            "training_phases": ["causal_lm_warmup", "mtp_horizon_training"],
            "total_parameters": count_parameters(model),
        }
        config_path = "./enhanced-qwen3-final/training_config.json"
        with open(config_path, "w") as f:
            json.dump(training_config, f, indent=2)
        print(f" Enhanced model, tokenizer, and config saved!")
    except Exception as e:
        print(f"Enhanced training failed with error: {e}")
        import traceback
        traceback.print_exc()
        return

    print("\nmTesting enhanced generation capabilities...")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    try:
        print("\n--- Standard Generation Test ---")
        input_ids = tokenizer("<s> [C]", return_tensors="pt").input_ids.to(device)
        with torch.no_grad():
            model.set_mtp_training(False)
            gen = model.generate(
                input_ids,
                max_length=25,
                top_k=50,
                top_p=0.9,
                temperature=1.0,
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                num_return_sequences=3,
            )
            for i, sequence in enumerate(gen):
                result = tokenizer.decode(sequence, skip_special_tokens=True)
                print(f"Generated SELFIES {i+1}: {result}")
        print("\n--- MTP Analysis Test ---")
        test_smiles = "[C]"
        test_input = tokenizer(test_smiles, return_tensors="pt", add_special_tokens=True).to(device)
        test_input = {k: v for k, v in test_input.items() if k != 'token_type_ids'}  # Remove token_type_ids
        with torch.no_grad():
            outputs = model(**test_input)
            if hasattr(model, 'mtp_head') and hasattr(model.mtp_head, 'prediction_heads'):
                hidden_states = model.model(test_input['input_ids']).last_hidden_state
                mtp_outputs = model.mtp_head(hidden_states)
                print(f"Input SELFIES: {test_smiles}")
                print(f"Tokenized: {tokenizer.convert_ids_to_tokens(test_input['input_ids'][0].tolist())}")
                for i, (key, logits) in enumerate(mtp_outputs.items()):
                    top_tokens = torch.topk(logits[0], k=3, dim=-1)
                    print(f"\n{key} predictions:")
                    for pos in range(min(5, logits.size(1))):
                        pos_preds = []
                        for j in range(3):
                            token_id = top_tokens.indices[pos, j].item()
                            prob = torch.softmax(logits[0, pos], dim=-1)[token_id].item()
                            token = tokenizer.id_to_token.get(token_id, '<UNK>')
                            pos_preds.append(f"{token}({prob:.3f})")
                        print(f"  Position {pos}: {', '.join(pos_preds)}")
        print("\nEnhanced generation tests completed!")
    except Exception as e:
        print(f"Enhanced generation test failed: {e}")
        import traceback
        traceback.print_exc()

    print("\nEnhanced Model Analysis:")
    print(f"Total parameters: {count_parameters(model):,}")
    mtp_params = sum(p.numel() for p in model.mtp_head.parameters() if p.requires_grad)
    horizon_params = sum(p.numel() for p in model.horizon_loss.parameters() if p.requires_grad)
    base_params = count_parameters(model) - mtp_params - horizon_params
    print(f"Base model parameters: {base_params:,}")
    print(f"MTP head parameters: {mtp_params:,}")
    print(f"Horizon loss parameters: {horizon_params:,}")
    print(f"Enhancement overhead: {((mtp_params + horizon_params) / base_params * 100):.2f}%")
    print(f"\n Enhanced Model Architecture:")
    print(f"- Base Model: Qwen2 with {config.num_hidden_layers} layers")  # Updated this line
    print(f"- Hidden Size: {config.hidden_size}")
    print(f"- Attention Heads: {config.num_attention_heads}")
    print(f"- Vocab Size: {config.vocab_size}")
    print(f"- MTP Future Tokens: {model.mtp_head.num_future_tokens}")
    print(f"- Horizon Loss Weights: Learnable")
    print(f"- Training Mode: {'MTP + Horizon Loss' if model.use_mtp_training else 'Standard Causal LM'}")
    print("\n Enhanced training pipeline completed successfully!")

if __name__ == "__main__":
    main()