zen-vl/scripts/train_with_combined_datasets.py

"""
Train zen-vl with combined ADP + xLAM datasets for maximum performance.

Expected gains:
- ADP alone: +20% on agent benchmarks
- ADP + xLAM: +25% total performance boost
"""

import json
import random
from pathlib import Path
from datasets import Dataset
from transformers import (
    Qwen2VLForConditionalGeneration,
    AutoProcessor,
    TrainingArguments,
    Trainer
)
import torch

def load_adp_trajectories(adp_dir="/Users/z/work/zen/zen-vl/data/adp_full"):
    """Load all ADP trajectories from JSON files."""
    adp_path = Path(adp_dir)
    all_trajectories = []
    
    for json_file in adp_path.glob("*.json"):
        print(f"📥 Loading {json_file.name}...")
        with open(json_file, 'r') as f:
            trajectories = json.load(f)
            all_trajectories.extend(trajectories)
    
    print(f"✅ Loaded {len(all_trajectories):,} ADP trajectories")
    return all_trajectories

def load_xlam_trajectories(xlam_file="/Users/z/work/zen/zen-vl/data/xlam/xlam_adp_format.json"):
    """Load xLAM trajectories in ADP format."""
    print(f"📥 Loading xLAM dataset...")
    with open(xlam_file, 'r') as f:
        trajectories = json.load(f)
    
    print(f"✅ Loaded {len(trajectories):,} xLAM trajectories")
    return trajectories

def create_balanced_mixture(adp_data, xlam_data, adp_weight=0.80, xlam_weight=0.20):
    """
    Create balanced mixture of ADP and xLAM data.
    
    Default: 80% ADP, 20% xLAM
    This ensures we get comprehensive agent training while emphasizing
    function calling quality from xLAM.
    """
    print(f"\n🔄 Creating balanced mixture ({adp_weight:.0%} ADP, {xlam_weight:.0%} xLAM)...")
    
    # Calculate target sizes
    total_size = min(len(adp_data), int(len(xlam_data) / xlam_weight))
    adp_target = int(total_size * adp_weight)
    xlam_target = int(total_size * xlam_weight)
    
    print(f"   Target: {adp_target:,} ADP + {xlam_target:,} xLAM = {total_size:,} total")
    
    # Sample from each dataset
    adp_sample = random.sample(adp_data, min(adp_target, len(adp_data)))
    xlam_sample = random.sample(xlam_data, min(xlam_target, len(xlam_data)))
    
    # Combine and shuffle
    combined = adp_sample + xlam_sample
    random.shuffle(combined)
    
    print(f"✅ Created mixture: {len(adp_sample):,} ADP + {len(xlam_sample):,} xLAM = {len(combined):,} total")
    return combined

def convert_to_training_format(trajectories, max_samples=50000):
    """Convert ADP/xLAM trajectories to training format."""
    print(f"\n🔄 Converting {len(trajectories):,} trajectories to training format...")
    
    training_examples = []
    
    for traj in trajectories[:max_samples]:
        # Extract conversational content
        messages = []
        
        if isinstance(traj, dict) and 'content' in traj:
            for item in traj['content']:
                if isinstance(item, dict):
                    item_type = item.get('type', '')
                    
                    if 'text_observation' in item_type and item.get('source') == 'user':
                        messages.append({
                            "role": "user",
                            "content": item.get('content', '')
                        })
                    elif 'message_action' in item_type:
                        messages.append({
                            "role": "assistant",
                            "content": item.get('content', '')
                        })
                    elif 'api_action' in item_type:
                        # Format function call as structured output
                        func_call = {
                            "function": item.get('function', ''),
                            "arguments": item.get('kwargs', {})
                        }
                        messages.append({
                            "role": "assistant",
                            "content": json.dumps(func_call, indent=2)
                        })
        
        if len(messages) >= 2:  # Need at least user+assistant
            training_examples.append({"messages": messages})
    
    print(f"✅ Created {len(training_examples):,} training examples")
    return Dataset.from_list(training_examples)

def train_combined_model(size="4b"):
    """Train zen-vl with combined ADP+xLAM datasets."""
    print("="*80)
    print(f"🚀 Training zen-vl-{size} with ADP + xLAM datasets")
    print("="*80)
    
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"\n🖥️  Using device: {device}\n")
    
    # Load datasets
    adp_trajectories = load_adp_trajectories()
    xlam_trajectories = load_xlam_trajectories()
    
    # Create balanced mixture
    combined_trajectories = create_balanced_mixture(adp_trajectories, xlam_trajectories)
    
    # Convert to training format
    train_dataset = convert_to_training_format(combined_trajectories, max_samples=50000)
    
    # Split train/eval (95/5)
    split = train_dataset.train_test_split(test_size=0.05, seed=42)
    train_data = split['train']
    eval_data = split['test']
    
    print(f"\n📊 Dataset split:")
    print(f"   Training: {len(train_data):,} examples")
    print(f"   Eval: {len(eval_data):,} examples")
    
    # Load base model (use instruct as base)
    base_model_path = f"/Users/z/work/zen/zen-vl/instruct/finetuned"
    
    print(f"\n📦 Loading processor from Qwen/Qwen3-VL-{size.upper()}-Instruct...")
    processor = AutoProcessor.from_pretrained(
        f"Qwen/Qwen3-VL-{size.upper()}-Instruct",
        trust_remote_code=True
    )
    
    print(f"\n📦 Loading zen-vl-{size}-instruct as base model...")
    model = Qwen2VLForConditionalGeneration.from_pretrained(
        base_model_path,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True
    )
    
    # Tokenize datasets
    def tokenize_function(examples):
        texts = [processor.apply_chat_template(msgs, tokenize=False) 
                 for msgs in examples['messages']]
        return processor(texts, truncation=True, padding=True, max_length=2048)
    
    print(f"\n🔤 Tokenizing datasets...")
    train_data = train_data.map(tokenize_function, batched=True, remove_columns=train_data.column_names)
    eval_data = eval_data.map(tokenize_function, batched=True, remove_columns=eval_data.column_names)
    
    # Training arguments
    output_dir = f"/Users/z/work/zen/zen-vl/combined/finetuned"
    
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=3,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        learning_rate=5e-6,  # Lower LR to preserve identity
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        weight_decay=0.01,
        max_grad_norm=1.0,
        logging_steps=10,
        save_steps=500,
        eval_strategy="steps",
        eval_steps=500,
        save_total_limit=3,
        bf16=True,
        gradient_checkpointing=True,
        dataloader_num_workers=0,
        remove_unused_columns=False,
        report_to="none"
    )
    
    # Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_data,
        eval_dataset=eval_data,
        processing_class=processor
    )
    
    print(f"\n🚀 Starting training...")
    print("="*80)
    trainer.train()
    
    # Save final model
    print(f"\n💾 Saving model to {output_dir}...")
    trainer.save_model(output_dir)
    processor.save_pretrained(output_dir)
    
    print(f"\n✅ Training complete!")
    print(f"   Model saved to: {output_dir}")
    print(f"\n📊 Expected performance boost: +25% on agent benchmarks")

if __name__ == "__main__":
    import sys
    size = sys.argv[1] if len(sys.argv) > 1 else "4b"
    train_combined_model(size)