prepare_xlam_data.py

#!/usr/bin/env python3
"""
Data preprocessing script for xLAM function calling dataset.
Downloads the xLAM dataset from Hugging Face and converts it to MLX-compatible JSONL format.
"""

import argparse
import json
import os
from pathlib import Path
from typing import Dict, Any

from datasets import load_dataset
from tqdm import tqdm


def process_xlam_sample(row: Dict[str, Any]) -> str:
    """
    Process a single xLAM dataset sample into function calling format.

    Format:
    <user>[user query]</user>

    <tools>
    [tool definitions]
    </tools>

    <calls>
    [expected function calls]
    </calls>

    Args:
        row: Dictionary containing 'query', 'tools', and 'answers' fields

    Returns:
        Formatted string for training
    """
    # Format user query
    formatted_query = f"<user>{row['query']}</user>\n\n"

    # Parse and format available tools
    try:
        parsed_tools = json.loads(row["tools"])
        tools_text = '\n'.join(str(tool) for tool in parsed_tools)
    except (json.JSONDecodeError, TypeError):
        tools_text = str(row["tools"])

    formatted_tools = f"<tools>{tools_text}</tools>\n\n"

    # Parse and format expected function calls
    try:
        parsed_answers = json.loads(row["answers"])
        answers_text = '\n'.join(str(answer) for answer in parsed_answers)
    except (json.JSONDecodeError, TypeError):
        answers_text = str(row["answers"])

    formatted_answers = f"<calls>{answers_text}</calls>"

    # Combine all parts (EOS token will be added by tokenizer during training)
    complete_text = formatted_query + formatted_tools + formatted_answers

    return complete_text


def download_and_process_xlam(
    output_dir: str = "data/xlam",
    train_samples: int = 800,
    valid_samples: int = 100,
    test_samples: int = 100,
    seed: int = 42
):
    """
    Download xLAM dataset and process it into train/valid/test splits.

    Args:
        output_dir: Directory to save the processed JSONL files
        train_samples: Number of training samples
        valid_samples: Number of validation samples
        test_samples: Number of test samples
        seed: Random seed for reproducible splits
    """
    print("=" * 80)
    print("xLAM Function Calling Dataset Preparation")
    print("=" * 80)

    # Create output directory
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    print(f"\n📁 Output directory: {output_path.absolute()}")
    print(f"📊 Sample sizes: train={train_samples}, valid={valid_samples}, test={test_samples}")
    print(f"🎲 Random seed: {seed}\n")

    # Load the xLAM dataset from Hugging Face
    print("📥 Downloading xLAM dataset from Hugging Face...")
    try:
        dataset = load_dataset("Salesforce/xlam-function-calling-60k", split="train")
        print(f"✅ Downloaded {len(dataset):,} samples from xLAM dataset\n")
    except Exception as e:
        print(f"❌ Error downloading dataset: {e}")
        print("\nTroubleshooting:")
        print("1. Check your internet connection")
        print("2. Make sure you have set your HuggingFace token if the dataset requires authentication")
        print("3. Run: huggingface-cli login")
        raise

    # Shuffle the dataset with fixed seed for reproducibility
    print(f"🔀 Shuffling dataset with seed {seed}...")
    dataset = dataset.shuffle(seed=seed)

    # Calculate total samples needed
    total_samples = train_samples + valid_samples + test_samples

    if total_samples > len(dataset):
        print(f"⚠️  Warning: Requested {total_samples} samples but dataset only has {len(dataset)}")
        print(f"    Using maximum available samples...")
        # Adjust proportionally
        ratio = len(dataset) / total_samples
        train_samples = int(train_samples * ratio)
        valid_samples = int(valid_samples * ratio)
        test_samples = len(dataset) - train_samples - valid_samples

    # Split the dataset
    splits = {
        "train": dataset.select(range(train_samples)),
        "valid": dataset.select(range(train_samples, train_samples + valid_samples)),
        "test": dataset.select(range(train_samples + valid_samples, train_samples + valid_samples + test_samples))
    }

    print(f"\n📊 Dataset splits:")
    print(f"   • Training:   {len(splits['train']):,} samples")
    print(f"   • Validation: {len(splits['valid']):,} samples")
    print(f"   • Test:       {len(splits['test']):,} samples")
    print(f"   • Total:      {sum(len(s) for s in splits.values()):,} samples\n")

    # Process and save each split
    for split_name, split_data in splits.items():
        output_file = output_path / f"{split_name}.jsonl"
        print(f"⚙️  Processing {split_name} split...")

        with open(output_file, "w") as f:
            for sample in tqdm(split_data, desc=f"Writing {split_name}.jsonl"):
                try:
                    processed_text = process_xlam_sample(sample)
                    json_line = {"text": processed_text}
                    f.write(json.dumps(json_line) + "\n")
                except Exception as e:
                    print(f"\n⚠️  Warning: Error processing sample: {e}")
                    continue

        print(f"✅ Saved {split_name}.jsonl ({len(split_data)} samples)\n")

    # Show a preview of the first training sample
    print("=" * 80)
    print("📋 Sample Preview (first training example):")
    print("=" * 80)

    with open(output_path / "train.jsonl", "r") as f:
        first_sample = json.loads(f.readline())
        sample_text = first_sample["text"]

        # Truncate if too long
        max_preview_length = 500
        if len(sample_text) > max_preview_length:
            preview = sample_text[:max_preview_length] + "\n... [truncated]"
        else:
            preview = sample_text

        print(preview)

    print("\n" + "=" * 80)
    print("✅ Dataset preparation complete!")
    print("=" * 80)
    print(f"\n📁 Files created in: {output_path.absolute()}")
    print("   • train.jsonl")
    print("   • valid.jsonl")
    print("   • test.jsonl")
    print("\n🚀 Next step: Run model conversion and training")


def main():
    parser = argparse.ArgumentParser(
        description="Download and preprocess xLAM dataset for MLX function calling training"
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="data/xlam",
        help="Directory to save processed JSONL files (default: data/xlam)"
    )
    parser.add_argument(
        "--train-samples",
        type=int,
        default=800,
        help="Number of training samples (default: 800)"
    )
    parser.add_argument(
        "--valid-samples",
        type=int,
        default=100,
        help="Number of validation samples (default: 100)"
    )
    parser.add_argument(
        "--test-samples",
        type=int,
        default=100,
        help="Number of test samples (default: 100)"
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for dataset shuffling (default: 42)"
    )

    args = parser.parse_args()

    download_and_process_xlam(
        output_dir=args.output_dir,
        train_samples=args.train_samples,
        valid_samples=args.valid_samples,
        test_samples=args.test_samples,
        seed=args.seed
    )


if __name__ == "__main__":
    main()