Create dataset.py

dataset.py

@@ -0,0 +1,763 @@
+"""
+This module provides PyTorch Dataset implementations for hierarchical VCF data
+"""
+
+import torch
+import json
+import pickle
+import logging
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional, Union, Any, Callable
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import pandas as pd
+
+from datasets import Dataset as HFDataset, DatasetDict
+from transformers import PreTrainedTokenizer
+
+from config import DataConfig, ModelConfig, ConfigManager
+from parser import VCFParser, MutationRecord
+from tokenizer import HierarchicalVCFTokenizer, HierarchicalDataCollator
+
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class HierarchicalVCFDataset(Dataset):
+    
+    def __init__(self,
+                 data_source: Union[str, Path, Dict, List],
+                 tokenizer: HierarchicalVCFTokenizer,
+                 config: Optional[DataConfig] = None,
+                 labels: Optional[Union[List, np.ndarray]] = None,
+                 transform: Optional[Callable] = None,
+                 target_transform: Optional[Callable] = None,
+                 cache_processed_data: bool = True):
+        """
+        Initialize the Hierarchical VCF Dataset.
+        Args:
+            data_source: Path to data file, or preprocessed data dict/list
+            tokenizer: Tokenizer for encoding mutations
+            config: Data configuration
+            labels: Optional labels for supervised learning
+            transform: Optional transform to apply to samples
+            target_transform: Optional transform to apply to labels
+            cache_processed_data: Whether to cache processed data
+        """
+        
+        self.config = config or DataConfig()
+        self.tokenizer = tokenizer
+        self.labels = labels
+        self.transform = transform
+        self.target_transform = target_transform
+        self.cache_processed_data = cache_processed_data
+        
+        # Load and process data
+        self.raw_data = self._load_data(data_source)
+        self.processed_data = self._process_data()
+        
+        # Validate data consistency
+        self._validate_data()
+        
+        # Dataset statistics
+        self.stats = self._compute_statistics()
+        
+        logger.info(f"Dataset initialized with {len(self.processed_data)} samples")
+        logger.info(f"Dataset statistics: {self.stats}")
+    
+    def _load_data(self, data_source: Union[str, Path, Dict, List]) -> Dict[str, Any]:
+        
+        if isinstance(data_source, (dict, list)):
+            # Data already loaded
+            if isinstance(data_source, list):
+                # Convert list to dict format
+                return {f"sample_{i}": sample for i, sample in enumerate(data_source)}
+            return data_source
+        
+        # Load from file
+        data_path = Path(data_source)
+        
+        if not data_path.exists():
+            raise FileNotFoundError(f"Data file not found: {data_path}")
+        
+        try:
+            if data_path.suffix.lower() == '.json':
+                with open(data_path, 'r') as f:
+                    return json.load(f)
+            
+            elif data_path.suffix.lower() == '.pkl':
+                with open(data_path, 'rb') as f:
+                    return pickle.load(f)
+            
+            elif data_path.suffix.lower() == '.vcf':
+                # Parse VCF file directly
+                parser = VCFParser(config=self.config)
+                return parser.parse_vcf_file(data_path)
+            
+            else:
+                raise ValueError(f"Unsupported file format: {data_path.suffix}")
+                
+        except Exception as e:
+            logger.error(f"Error loading data from {data_path}: {e}")
+            raise
+    
+    def _process_data(self) -> List[Dict[str, Any]]:
+        """Raw hierarchical data into dataset format."""
+        
+        processed_samples = []
+        
+        for sample_id, sample_data in self.raw_data.items():
+            try:
+                # Convert to standard format if needed
+                standardized_sample = self._standardize_sample_format(sample_data)
+                
+                # Filter samples based on configuration
+                if self._should_include_sample(standardized_sample):
+                    # Encode the sample
+                    encoded_sample = self.tokenizer.encode_hierarchical_sample(standardized_sample)
+                    
+                    processed_sample = {
+                        'sample_id': sample_id,
+                        'encoded_data': encoded_sample,
+                        'raw_data': standardized_sample if not self.cache_processed_data else None
+                    }
+                    
+                    processed_samples.append(processed_sample)
+                
+            except Exception as e:
+                logger.warning(f"Error processing sample {sample_id}: {e}")
+                continue
+        
+        return processed_samples
+    
+    def _standardize_sample_format(self, sample_data: Dict[str, Any]) -> Dict[str, Any]:
+        
+        # Handle different input formats
+        if 'mutations' in sample_data:
+            # Format: {'mutations': [...]}
+            return self._convert_flat_to_hierarchical(sample_data['mutations'])
+        
+        elif isinstance(sample_data, dict) and all(
+            isinstance(v, dict) for v in sample_data.values()
+        ):
+            # Already in hierarchical format
+            return sample_data
+        
+        else:
+            # Assume it's a list of mutations
+            return self._convert_flat_to_hierarchical(sample_data)
+    
+    def _convert_flat_to_hierarchical(self, mutations: List[Dict]) -> Dict[str, Any]:
+        """Convert flat mutation list to hierarchical format."""
+        
+        hierarchical = {}
+        
+        for mutation in mutations:
+            # Extract hierarchical keys
+            pathway = mutation.get('pathway', 'Unknown_Pathway')
+            chromosome = mutation.get('chromosome', mutation.get('chrom', 'Unknown'))  
+            gene = mutation.get('gene', mutation.get('gene_id', 'Unknown_Gene'))
+            
+            # Initialize nested structure
+            if pathway not in hierarchical:
+                hierarchical[pathway] = {}
+            if chromosome not in hierarchical[pathway]:
+                hierarchical[pathway][chromosome] = {}
+            if gene not in hierarchical[pathway][chromosome]:
+                hierarchical[pathway][chromosome][gene] = []
+            
+            # Add mutation
+            hierarchical[pathway][chromosome][gene].append(mutation)
+        
+        return hierarchical
+    
+    def _should_include_sample(self, sample_data: Dict[str, Any]) -> bool:
+        """Determine if sample should be included based on filtering criteria."""
+        
+        # Count total mutations
+        total_mutations = 0
+        for pathway_data in sample_data.values():
+            for chrom_data in pathway_data.values():
+                for gene_mutations in chrom_data.values():
+                    total_mutations += len(gene_mutations)
+        
+        # Apply filters
+        if total_mutations < self.config.min_mutations_per_sample:
+            return False
+        
+        if total_mutations > self.config.max_mutations_per_sample:
+            return False
+        
+        return True
+    
+    def _validate_data(self) -> None:
+        
+        if len(self.processed_data) == 0:
+            raise ValueError("No valid samples found in dataset")
+        
+        if self.labels is not None:
+            if len(self.labels) != len(self.processed_data):
+                raise ValueError(
+                    f"Number of labels ({len(self.labels)}) doesn't match "
+                    f"number of samples ({len(self.processed_data)})"
+                )
+    
+    def _compute_statistics(self) -> Dict[str, Any]:
+        """CDataset statistics."""
+        
+        stats = {
+            'num_samples': len(self.processed_data),
+            'num_pathways': set(),
+            'num_chromosomes': set(),
+            'num_genes': set(),
+            'mutations_per_sample': [],
+            'genes_per_sample': [],
+            'pathways_per_sample': []
+        }
+        
+        for sample in self.processed_data:
+            encoded_data = sample['encoded_data']
+            
+            sample_pathways = len(encoded_data)
+            sample_genes = 0
+            sample_mutations = 0
+            
+            for pathway_token, chromosomes in encoded_data.items():
+                stats['num_pathways'].add(pathway_token)
+                
+                for chrom_token, genes in chromosomes.items():
+                    stats['num_chromosomes'].add(chrom_token)
+                    
+                    for gene_token, mutations in genes.items():
+                        stats['num_genes'].add(gene_token)
+                        sample_genes += 1
+                        
+                        # Count mutations (assuming 'impact' field exists)
+                        if 'impact' in mutations:
+                            sample_mutations += len(mutations['impact'])
+            
+            stats['mutations_per_sample'].append(sample_mutations)
+            stats['genes_per_sample'].append(sample_genes)
+            stats['pathways_per_sample'].append(sample_pathways)
+        
+        # Convert sets to counts
+        stats['unique_pathways'] = len(stats['num_pathways'])
+        stats['unique_chromosomes'] = len(stats['num_chromosomes'])
+        stats['unique_genes'] = len(stats['num_genes'])
+        
+        # Compute summary statistics
+        if stats['mutations_per_sample']:
+            stats['avg_mutations_per_sample'] = np.mean(stats['mutations_per_sample'])
+            stats['std_mutations_per_sample'] = np.std(stats['mutations_per_sample'])
+        
+        if stats['genes_per_sample']:
+            stats['avg_genes_per_sample'] = np.mean(stats['genes_per_sample'])
+            stats['std_genes_per_sample'] = np.std(stats['genes_per_sample'])
+        
+        # Remove raw sets
+        del stats['num_pathways'], stats['num_chromosomes'], stats['num_genes']
+        
+        return stats
+    
+    def __len__(self) -> int:
+        """Number of samples in the dataset."""
+        return len(self.processed_data)
+    
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        """Single sample from the dataset."""
+        
+        if idx >= len(self.processed_data):
+            raise IndexError(f"Index {idx} out of range for dataset of size {len(self)}")
+        
+        sample = self.processed_data[idx].copy()
+        
+        # Apply transforms
+        if self.transform:
+            sample['encoded_data'] = self.transform(sample['encoded_data'])
+        
+        # Add label if available
+        if self.labels is not None:
+            label = self.labels[idx]
+            if self.target_transform:
+                label = self.target_transform(label)
+            sample['label'] = label
+        
+        return sample
+    
+    def get_sample_by_id(self, sample_id: str) -> Optional[Dict[str, Any]]:
+        for i, sample in enumerate(self.processed_data):
+            if sample['sample_id'] == sample_id:
+                return self.__getitem__(i)
+        return None
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        return self.stats.copy()
+    
+    def save_dataset(self, save_path: Union[str, Path], format: str = 'pickle') -> None:
+        """
+        Args:
+            save_path: Path to save the dataset
+            format: Save format ('pickle', 'json')
+        """
+        save_path = Path(save_path)
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        dataset_info = {
+            'processed_data': self.processed_data,
+            'labels': self.labels.tolist() if isinstance(self.labels, np.ndarray) else self.labels,
+            'stats': self.stats,
+            'config': self.config.__dict__ if hasattr(self.config, '__dict__') else None
+        }
+        
+        if format.lower() == 'pickle':
+            with open(save_path, 'wb') as f:
+                pickle.dump(dataset_info, f)
+        
+        elif format.lower() == 'json':
+            with open(save_path, 'w') as f:
+                json.dump(dataset_info, f, indent=2, default=str)
+        
+        else:
+            raise ValueError(f"Unsupported save format: {format}")
+        
+        logger.info(f"Dataset saved to {save_path}")
+    
+    @classmethod
+    def load_dataset(cls, 
+                     load_path: Union[str, Path], 
+                     tokenizer: HierarchicalVCFTokenizer,
+                     format: str = 'auto') -> 'HierarchicalVCFDataset':
+        """
+        Args:
+            load_path: Path to load the dataset from
+            tokenizer: Tokenizer instance
+            format: Load format ('pickle', 'json', 'auto')
+            
+        Returns:
+            Loaded dataset instance
+        """
+        load_path = Path(load_path)
+        
+        if not load_path.exists():
+            raise FileNotFoundError(f"Dataset file not found: {load_path}")
+        
+        # Determine format
+        if format == 'auto':
+            format = 'pickle' if load_path.suffix == '.pkl' else 'json'
+        
+        # Load data
+        if format.lower() == 'pickle':
+            with open(load_path, 'rb') as f:
+                dataset_info = pickle.load(f)
+        
+        elif format.lower() == 'json':
+            with open(load_path, 'r') as f:
+                dataset_info = json.load(f)
+        
+        else:
+            raise ValueError(f"Unsupported load format: {format}")
+        
+        # Create dataset instance
+        dataset = cls.__new__(cls)
+        dataset.tokenizer = tokenizer
+        dataset.processed_data = dataset_info['processed_data']
+        dataset.labels = dataset_info.get('labels')
+        dataset.stats = dataset_info.get('stats', {})
+        dataset.config = dataset_info.get('config', DataConfig())
+        dataset.transform = None
+        dataset.target_transform = None
+        dataset.cache_processed_data = True
+        
+        return dataset
+
+
+class HierarchicalVCFDataModule:
+    """
+    Manage train/validation/test splits of hierarchical VCF data.
+    """
+    
+    def __init__(self,
+                 data_source: Union[str, Path, Dict],
+                 tokenizer: HierarchicalVCFTokenizer,
+                 config: Optional[DataConfig] = None,
+                 labels: Optional[Union[List, np.ndarray]] = None,
+                 train_split: float = 0.8,
+                 val_split: float = 0.1,
+                 test_split: float = 0.1,
+                 stratify: bool = True,
+                 random_seed: int = 42):
+        """
+        Args:
+            data_source: Source of the data
+            tokenizer: Tokenizer for encoding
+            config: Data configuration
+            labels: Labels for supervised learning
+            train_split: Proportion for training
+            val_split: Proportion for validation
+            test_split: Proportion for testing
+            stratify: Whether to stratify splits by labels
+            random_seed: Random seed for reproducibility
+        """
+        
+        self.config = config or DataConfig()
+        self.tokenizer = tokenizer
+        self.train_split = train_split
+        self.val_split = val_split
+        self.test_split = test_split
+        self.stratify = stratify
+        self.random_seed = random_seed
+        
+        # Validate splits
+        if abs(train_split + val_split + test_split - 1.0) > 1e-6:
+            raise ValueError("Train, validation, and test splits must sum to 1.0")
+        
+        # Load full dataset
+        self.full_dataset = HierarchicalVCFDataset(
+            data_source=data_source,
+            tokenizer=tokenizer,
+            config=config,
+            labels=labels
+        )
+        
+        # Create splits
+        self.train_dataset, self.val_dataset, self.test_dataset = self._create_splits()
+        
+        logger.info(f"Data module initialized:")
+        logger.info(f"  Train: {len(self.train_dataset)} samples")
+        logger.info(f"  Validation: {len(self.val_dataset)} samples") 
+        logger.info(f"  Test: {len(self.test_dataset)} samples")
+    
+    def _create_splits(self) -> Tuple[Dataset, Dataset, Dataset]:
+        
+        np.random.seed(self.random_seed)
+        
+        indices = np.arange(len(self.full_dataset))
+        
+        if self.stratify and self.full_dataset.labels is not None:
+            # Stratified split
+            from sklearn.model_selection import train_test_split
+            
+            # First split: train vs (val + test)
+            train_idx, temp_idx = train_test_split(
+                indices,
+                test_size=(self.val_split + self.test_split),
+                stratify=[self.full_dataset.labels[i] for i in indices],
+                random_state=self.random_seed
+            )
+            
+            # Second split: val vs test
+            if self.test_split > 0:
+                val_idx, test_idx = train_test_split(
+                    temp_idx,
+                    test_size=self.test_split / (self.val_split + self.test_split),
+                    stratify=[self.full_dataset.labels[i] for i in temp_idx],
+                    random_state=self.random_seed
+                )
+            else:
+                val_idx = temp_idx
+                test_idx = np.array([])
+            
+        else:
+            # Random split
+            np.random.shuffle(indices)
+            
+            train_end = int(self.train_split * len(indices))
+            val_end = int((self.train_split + self.val_split) * len(indices))
+            
+            train_idx = indices[:train_end]
+            val_idx = indices[train_end:val_end]
+            test_idx = indices[val_end:]
+        
+        # Create subset datasets
+        train_dataset = self._create_subset(train_idx)
+        val_dataset = self._create_subset(val_idx)
+        test_dataset = self._create_subset(test_idx)
+        
+        return train_dataset, val_dataset, test_dataset
+    
+    def _create_subset(self, indices: np.ndarray) -> Dataset:
+        """Create a subset dataset from indices."""
+        
+        subset_data = [self.full_dataset.processed_data[i] for i in indices]
+        subset_labels = None
+        
+        if self.full_dataset.labels is not None:
+            if isinstance(self.full_dataset.labels, np.ndarray):
+                subset_labels = self.full_dataset.labels[indices]
+            else:
+                subset_labels = [self.full_dataset.labels[i] for i in indices]
+        
+        # Create new dataset instance
+        dataset = HierarchicalVCFDataset.__new__(HierarchicalVCFDataset)
+        dataset.tokenizer = self.tokenizer
+        dataset.config = self.config
+        dataset.processed_data = subset_data
+        dataset.labels = subset_labels
+        dataset.transform = None
+        dataset.target_transform = None
+        dataset.cache_processed_data = True
+        dataset.stats = dataset._compute_statistics()
+        
+        return dataset
+    
+    def get_dataloaders(self, 
+                       batch_size: int = 16,
+                       num_workers: int = 0,
+                       collate_fn: Optional[Callable] = None) -> Tuple[DataLoader, DataLoader, DataLoader]:
+        """
+        Args:
+            batch_size: Batch size for data loading
+            num_workers: Number of worker processes
+            collate_fn: Custom collate function
+            
+        Returns:
+            Tuple of (train_loader, val_loader, test_loader)
+        """
+        
+        if collate_fn is None:
+            collate_fn = HierarchicalDataCollator(self.tokenizer)
+        
+        train_loader = DataLoader(
+            self.train_dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            collate_fn=collate_fn
+        )
+        
+        val_loader = DataLoader(
+            self.val_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            collate_fn=collate_fn
+        )
+        
+        test_loader = DataLoader(
+            self.test_dataset,
+            batch_size=batch_size,
+            shuffle=False,
+            num_workers=num_workers,
+            collate_fn=collate_fn
+        )
+        
+        return train_loader, val_loader, test_loader
+
+
+class HuggingFaceDatasetAdapter:
+    """
+    Convert hierarchical VCF data to Hugging Face Dataset format.
+    """
+    
+    def __init__(self, vcf_dataset: HierarchicalVCFDataset):
+        self.vcf_dataset = vcf_dataset
+    
+    def to_huggingface_dataset(self) -> DatasetDict:
+        """
+        Returns:
+            HuggingFace DatasetDict
+        """
+        
+        # Flatten hierarchical data for HF compatibility
+        flattened_data = []
+        
+        for sample in self.vcf_dataset.processed_data:
+            sample_id = sample['sample_id']
+            encoded_data = sample['encoded_data']
+            
+            # Convert hierarchical structure to flattened format
+            flattened_sample = {
+                'sample_id': sample_id,
+                'pathways': list(encoded_data.keys()),
+                'num_pathways': len(encoded_data),
+                'encoded_mutations': self._flatten_mutations(encoded_data)
+            }
+            
+            flattened_data.append(flattened_sample)
+        
+        # Add labels if available
+        if self.vcf_dataset.labels is not None:
+            for i, sample in enumerate(flattened_data):
+                sample['label'] = self.vcf_dataset.labels[i]
+        
+        # Create HuggingFace dataset
+        hf_dataset = HFDataset.from_list(flattened_data)
+        
+        return DatasetDict({'train': hf_dataset})
+    
+    def _flatten_mutations(self, encoded_data: Dict) -> Dict[str, List]:
+        """Flatten hierarchical mutations for HF compatibility."""
+        
+        all_impacts = []
+        all_refs = []
+        all_alts = []
+        
+        for pathway_token, chromosomes in encoded_data.items():
+            for chrom_token, genes in chromosomes.items():
+                for gene_token, mutations in genes.items():
+                    if 'impact' in mutations:
+                        all_impacts.extend(mutations['impact'])
+                    if 'ref' in mutations:
+                        all_refs.extend(mutations['ref'])
+                    if 'alt' in mutations:
+                        all_alts.extend(mutations['alt'])
+        
+        return {
+            'impacts': all_impacts,
+            'refs': all_refs,
+            'alts': all_alts
+        }
+
+
+def create_dataset_from_config(config_manager: ConfigManager,
+                              tokenizer: HierarchicalVCFTokenizer,
+                              labels: Optional[List] = None) -> HierarchicalVCFDataset:
+    
+    data_config = config_manager.data_config
+    
+    if not data_config.vcf_file_path:
+        raise ValueError("VCF file path not specified in configuration")
+    
+    return HierarchicalVCFDataset(
+        data_source=data_config.vcf_file_path,
+        tokenizer=tokenizer,
+        config=data_config,
+        labels=labels
+    )
+
+
+def create_data_module_from_config(config_manager: ConfigManager,
+                                  tokenizer: HierarchicalVCFTokenizer,
+                                  labels: Optional[List] = None) -> HierarchicalVCFDataModule:
+    
+    data_config = config_manager.data_config
+    
+    if not data_config.vcf_file_path:
+        raise ValueError("VCF file path not specified in configuration")
+    
+    return HierarchicalVCFDataModule(
+        data_source=data_config.vcf_file_path,
+        tokenizer=tokenizer,
+        config=data_config,
+        labels=labels
+    )
+
+
+# Utility functions for data preprocessing
+def create_synthetic_labels(dataset: HierarchicalVCFDataset, 
+                           label_type: str = 'random',
+                           num_classes: int = 2) -> np.ndarray:
+    """
+    Create synthetic labels for testing purposes.
+    
+    Args:
+        dataset: VCF dataset
+        label_type: Type of labels ('random', 'mutation_count_based')
+        num_classes: Number of classes for classification
+        
+    Returns:
+        Array of synthetic labels
+    """
+    
+    num_samples = len(dataset)
+    
+    if label_type == 'random':
+        return np.random.randint(0, num_classes, size=num_samples)
+    
+    elif label_type == 'mutation_count_based':
+        # Create labels based on mutation count thresholds
+        mutation_counts = dataset.stats['mutations_per_sample']
+        threshold = np.median(mutation_counts)
+        
+        labels = []
+        for count in mutation_counts:
+            if num_classes == 2:
+                labels.append(1 if count > threshold else 0)
+            else:
+                # Divide into quantiles
+                percentiles = np.linspace(0, 100, num_classes + 1)
+                thresholds = np.percentile(mutation_counts, percentiles[1:-1])
+                
+                label = 0
+                for i, t in enumerate(thresholds):
+                    if count > t:
+                        label = i + 1
+                    else:
+                        break
+                labels.append(label)
+        
+        return np.array(labels)
+    
+    else:
+        raise ValueError(f"Unknown label type: {label_type}")
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    from tokenizer import create_tokenizer_from_config
+    
+    # Example usage
+    config_manager = ConfigManager()
+    config_manager.data_config.vcf_file_path = "example_data.json"
+    
+    # Create tokenizer
+    tokenizer = create_tokenizer_from_config(config_manager)
+    
+    # Example data
+    example_data = {
+        'sample1': {
+            'pathway1': {
+                'chr1': {
+                    'gene1': [
+                        {'impact': 'HIGH', 'reference': 'A', 'alternate': 'T'},
+                        {'impact': 'MODERATE', 'reference': 'G', 'alternate': 'C'}
+                    ]
+                }
+            }
+        },
+        'sample2': {
+            'pathway2': {
+                'chr2': {
+                    'gene2': [
+                        {'impact': 'LOW', 'reference': 'T', 'alternate': 'A'}
+                    ]
+                }
+            }
+        }
+    }
+    
+    # Build tokenizer vocabulary
+    tokenizer.build_vocabulary(example_data)
+    
+    # Create dataset
+    dataset = HierarchicalVCFDataset(
+        data_source=example_data,
+        tokenizer=tokenizer
+    )
+    
+    # Create synthetic labels
+    labels = create_synthetic_labels(dataset, label_type='random', num_classes=2)
+    dataset.labels = labels
+    
+    # Create data module
+    data_module = HierarchicalVCFDataModule(
+        data_source=example_data,
+        tokenizer=tokenizer,
+        labels=labels,
+        train_split=0.6,
+        val_split=0.2,
+        test_split=0.2
+    )
+    
+    # Get data loaders
+    train_loader, val_loader, test_loader = data_module.get_dataloaders(batch_size=2)
+    
+    # Test data loading
+    for batch in train_loader:
+        print(f"Batch size: {batch['batch_size']}")
+        print(f"Sample IDs: {[s.get('sample_id', 'N/A') for s in batch['samples']]}")
+        break
+    
+    print(f"Dataset statistics: {dataset.get_statistics()}")