[data] Expand dataset building for easier customization

The current dataset builders are restricted to configurations the framework directly knows about: https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/76125731c71ebc21462a63274c00cc2114377612/src/megatron/bridge/data/utils.py#L150-L169. This covers pretraining with megatron core and the basics of finetuning, with some customizations for HF datasets.

This design forces users to either:

• Modify framework source code to add new dataset types
• Plug their custom datasets into existing configurations

However, custom datasets may not clearly fit into these patterns. The framework cannot have builders for each possible custom dataset within its source. Instead, the framework should offer a protocol to handle instantiating datasets

Proposal:

Define an interface for how users can provide custom datasets given inputs from the framework during the setup

# In src/megatron/bridge/data/protocols.py
from abc import ABC, abstractmethod
from typing import Optional, Any
from dataclasses import dataclass
from megatron.bridge.training.config import DataloaderConfig
from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer

@dataclass(frozen=True)
class DatasetBuildContext:
    """Clean interface that encapsulates framework internals."""
    train_samples: int
    valid_samples: int  
    test_samples: int
    tokenizer: Optional[MegatronTokenizer] = None

@dataclass
class DatasetProvider(DataloaderConfig, ABC):
    """Abstract base class for custom dataset configurations.
    
    Provides a clean interface for users to implement their own dataset builders
    while automatically inheriting all DataloaderConfig functionality.
    
    Users must:
    1. Inherit from this class
    2. Implement the build_datasets() method
    3. Use @dataclass decorator (already provided by inheritance)
    
    Example:
        @dataclass
        class S3DatasetConfig(CustomDatasetConfig):
            bucket_name: str
            data_prefix: str
            
            def build_datasets(self, context: DatasetBuildContext):
                # implementation
                return train_ds, valid_ds, test_ds
    """
    
    @abstractmethod
    def build_datasets(self, context: DatasetBuildContext) -> tuple[Optional[Any], Optional[Any], Optional[Any]]:
        """Build train, validation, and test datasets.
        
        Args:
            context: Build context with sample counts and tokenizer
            
        Returns:
            Tuple of (train_dataset, valid_dataset, test_dataset)
            Any element can be None if that split shouldn't be created.
        """
        pass

Leverage this to extend the hardcoded registry

def get_dataset_provider(dataset_config) -> Callable:
    """Get provider function, supporting both legacy registry and new protocol."""
    
    # Check if config implements the DatasetProvider protocol
    if isinstance(dataset_config, DatasetProvider):
        
        def protocol_adapter(train_val_test_num_samples: list[int], config, tokenizer=None):
            context = DatasetBuildContext(
                train_samples=train_val_test_num_samples[0],
                valid_samples=train_val_test_num_samples[1], 
                test_samples=train_val_test_num_samples[2],
                tokenizer=tokenizer
            )
            return config.build_datasets(context)
            
        return protocol_adapter
    
    # Fall back to existing registry
    return _REGISTRY[type(dataset_config)]

And include this interface in the typing for the dataset config here