Fix the Data issue and Add the logging setups for better Experince.

Author: codewithdark-git

quantllm/__init__.py

@@ -11,15 +11,25 @@
     TrainingLogger
 )
 from .hub import HubManager, CheckpointManager
+from .utils.optimizations import get_optimal_training_settings
+from .utils.log_config import configure_logging, enable_logging
 
 from .config import (
     ModelConfig,
     DatasetConfig,
     TrainingConfig
 )
 
+# Configure package-wide logging
+configure_logging()
+
 __version__ = "0.1.0"
 
+# Package metadata
+__title__ = "QuantLLM"
+__description__ = "Efficient Quantized LLM Fine-Tuning Library"
+__author__ = "QuantLLM Team"
+
 __all__ = [
     # Model
     "Model",
@@ -42,5 +52,18 @@
     # Configuration
     "ModelConfig",
     "DatasetConfig",
-    "TrainingConfig"
-] 
+    "TrainingConfig",
+    
+    # Utilities
+    "get_optimal_training_settings",
+    "configure_logging",
+    "enable_logging",
+]
+
+# Initialize package-level logger with fancy welcome message
+logger = TrainingLogger()
+logger.log_success(f"""
+✨ QuantLLM v{__version__} initialized successfully ✨
+🚀 Efficient Quantized Language Model Fine-Tuning
+📚 Documentation: https://github.com/yourusername/QuantLLM
+""")

quantllm/data/dataset_preprocessor.py

@@ -1,26 +1,33 @@
 from datasets import Dataset
-from typing import Optional, Dict, Any, Callable
+from typing import Optional, Dict, Any, Callable, Tuple
 from transformers import PreTrainedTokenizer
 from ..trainer.logger import TrainingLogger
+from tqdm.auto import tqdm
+import logging
+import warnings
+
+# Disable unnecessary logging
+logging.getLogger("tokenizers").setLevel(logging.ERROR)
+warnings.filterwarnings("ignore")
 
 class DatasetPreprocessor:
-    def __init__(self, tokenizer, logger=None):
+    def __init__(self, tokenizer: PreTrainedTokenizer, logger: Optional[TrainingLogger] = None):
         self.tokenizer = tokenizer
         self.logger = logger or TrainingLogger()
 
         # Set pad token if not set
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
             self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-            print("Set pad token to eos token")
+            self.logger.log_info("Set pad token to eos token")
 
     def validate_datasets(self, datasets):
         """Validate input datasets."""
         for dataset in datasets:
             if dataset is not None and not isinstance(dataset, Dataset):
                 raise ValueError(f"Expected Dataset object, got {type(dataset)}")
 
-    def preprocess_text(self, text):
+    def preprocess_text(self, text: str) -> str:
         """Basic text preprocessing"""
         if not text:
             return ""
@@ -30,98 +37,98 @@ def preprocess_text(self, text):
 
     def tokenize_dataset(
         self,
-        train_dataset,
-        val_dataset=None,
-        test_dataset=None,
+        train_dataset: Dataset,
+        val_dataset: Optional[Dataset] = None,
+        test_dataset: Optional[Dataset] = None,
         max_length: int = 512,
         text_column: str = "text",
-        label_column: str = None,
+        label_column: Optional[str] = None,
         batch_size: int = 1000
-    ):
-        """Tokenize datasets with preprocessing."""
+    ) -> Tuple[Dataset, Optional[Dataset], Optional[Dataset]]:
+        """Tokenize datasets with preprocessing and progress bars."""
         try:
             self.validate_datasets([train_dataset, val_dataset, test_dataset])
 
             def process_and_tokenize_batch(examples):
-                # Get texts and preprocess
+                # Get texts and preprocess with progress indication
                 texts = examples[text_column]
                 if not isinstance(texts, list):
                     texts = [texts]
+                
+                # Preprocess texts
                 texts = [self.preprocess_text(text) for text in texts]
 
                 try:
                     # Tokenize with padding and truncation
-                    # Use max_length + 1 to account for the shift we'll do later
                     tokenized = self.tokenizer(
                         texts,
                         padding="max_length",
                         truncation=True,
-                        max_length=max_length + 1,  # Add 1 to account for shift
+                        max_length=max_length + 1,  # Add 1 for shift
                         return_tensors=None
                     )
 
+                    # For causal language modeling, prepare shifted sequences
                     input_ids = tokenized["input_ids"]
                     attention_mask = tokenized["attention_mask"]
 
-                    # Now shift to create inputs and labels
-                    # inputs will be [:-1] and labels will be [1:]
+                    # Prepare shifted sequences for input and labels
                     labels = [ids[1:] for ids in input_ids]
                     input_ids = [ids[:-1] for ids in input_ids]
                     attention_mask = [mask[:-1] for mask in attention_mask]
 
-                    # Verify all sequences have the expected length
+                    # Verify sequence lengths
                     expected_length = max_length
-                    if not all(len(seq) == expected_length for seq in input_ids):
-                        raise ValueError(f"Input sequence lengths don't match. Expected {expected_length}")
-                    if not all(len(seq) == expected_length for seq in attention_mask):
-                        raise ValueError(f"Attention mask lengths don't match. Expected {expected_length}")
-                    if not all(len(seq) == expected_length for seq in labels):
-                        raise ValueError(f"Label sequence lengths don't match. Expected {expected_length}")
+                    assert all(len(seq) == expected_length for seq in input_ids), "Input sequence lengths don't match"
+                    assert all(len(seq) == expected_length for seq in attention_mask), "Attention mask lengths don't match"
+                    assert all(len(seq) == expected_length for seq in labels), "Label sequence lengths don't match"
 
                     result = {
                         "input_ids": input_ids,
                         "attention_mask": attention_mask,
                         "labels": labels
                     }
 
-                    self.logger.log_info(f"Tokenized batch of {len(texts)} texts")
                     return result
 
                 except Exception as e:
                     self.logger.log_error(f"Error tokenizing batch: {str(e)}")
                     raise
 
-            # Process datasets
+            # Process datasets with overall progress bars
+            self.logger.log_info("Processing training dataset")
             train_tokenized = train_dataset.map(
                 process_and_tokenize_batch,
                 batched=True,
                 batch_size=batch_size,
                 remove_columns=train_dataset.column_names,
                 desc="Tokenizing training set"
             )
-            self.logger.log_info(f"Tokenized training dataset: {len(train_tokenized)} examples")
+            self.logger.log_success(f"Tokenized training dataset: {len(train_tokenized)} examples")
 
             val_tokenized = None
             if val_dataset is not None:
+                self.logger.log_info("Processing validation dataset")
                 val_tokenized = val_dataset.map(
                     process_and_tokenize_batch,
                     batched=True,
                     batch_size=batch_size,
                     remove_columns=val_dataset.column_names,
                     desc="Tokenizing validation set"
                 )
-                self.logger.log_info(f"Tokenized validation dataset: {len(val_tokenized)} examples")
+                self.logger.log_success(f"Tokenized validation dataset: {len(val_tokenized)} examples")
 
             test_tokenized = None
             if test_dataset is not None:
+                self.logger.log_info("Processing test dataset")
                 test_tokenized = test_dataset.map(
                     process_and_tokenize_batch,
                     batched=True,
                     batch_size=batch_size,
                     remove_columns=test_dataset.column_names,
                     desc="Tokenizing test set"
                 )
-                self.logger.log_info(f"Tokenized test dataset: {len(test_tokenized)} examples")
+                self.logger.log_success(f"Tokenized test dataset: {len(test_tokenized)} examples")
 
             # Set format to PyTorch tensors
             train_tokenized.set_format("torch")

quantllm/data/dataset_splitter.py

@@ -1,11 +1,26 @@
-from datasets import Dataset
-from typing import Optional, Dict, Any, Tuple
+from datasets import Dataset, DatasetDict
+from typing import Optional, Tuple, Union
+import numpy as np
 from ..trainer.logger import TrainingLogger
+from tqdm.auto import tqdm
+import logging
+
+# Configure logging
+logging.getLogger("datasets").setLevel(logging.WARNING)
 
 class DatasetSplitter:
-    def __init__(self, logger=None):
+    def __init__(self, logger: Optional[TrainingLogger] = None):
+        """Initialize dataset splitter."""
         self.logger = logger or TrainingLogger()
-        
+
+    def _get_dataset_from_dict(self, dataset: Union[Dataset, DatasetDict], split: str = "train") -> Dataset:
+        """Extract dataset from DatasetDict if needed."""
+        if isinstance(dataset, DatasetDict):
+            if split in dataset:
+                return dataset[split]
+            raise ValueError(f"DatasetDict does not contain split '{split}'")
+        return dataset
+
     def validate_split_params(self, train_size: float, val_size: float, test_size: float = None):
         """Validate split parameters."""
         if train_size <= 0 or train_size >= 1:
@@ -55,52 +70,104 @@ def train_test_split(
             self.logger.log_error(f"Error splitting dataset: {str(e)}")
             raise
 
-    def train_val_test_split(self, dataset, train_size: float, val_size: float, test_size: float = None):
-        """Split dataset into train, validation and test sets."""
+    def train_val_test_split(
+        self,
+        dataset: Union[Dataset, DatasetDict],
+        train_size: float = 0.8,
+        val_size: float = 0.1,
+        test_size: float = 0.1,
+        shuffle: bool = True,
+        seed: int = 42,
+        split: str = "train"
+    ) -> Tuple[Dataset, Dataset, Dataset]:
+        """
+        Split dataset into train, validation and test sets with progress indication.
+        
+        Args:
+            dataset (Dataset or DatasetDict): Dataset to split
+            train_size (float): Proportion of training set
+            val_size (float): Proportion of validation set
+            test_size (float): Proportion of test set
+            shuffle (bool): Whether to shuffle the dataset
+            seed (int): Random seed
+            split (str): Which split to use if dataset is a DatasetDict
+            
+        Returns:
+            Tuple[Dataset, Dataset, Dataset]: Train, validation and test datasets
+        """
         try:
-            if not isinstance(dataset, Dataset):
-                if isinstance(dataset, dict) and 'train' in dataset:
-                    dataset = dataset['train']
-                else:
-                    raise ValueError(f"Expected Dataset object or dict with 'train' key, got {type(dataset)}")
-                    
-            if test_size is None:
-                test_size = 1.0 - train_size - val_size
-                
-            self.validate_split_params(train_size, val_size, test_size)
+            # Get the actual dataset if we have a DatasetDict
+            dataset = self._get_dataset_from_dict(dataset, split)
 
-            # If dataset is already split
-            if isinstance(dataset, dict) and all(k in dataset for k in ['train', 'validation', 'test']):
-                self.logger.log_info("Dataset already contains train/validation/test splits")
-                return dataset['train'], dataset['validation'], dataset['test']
+            # Validate split proportions
+            total = train_size + val_size + test_size
+            if not np.isclose(total, 1.0):
+                raise ValueError(f"Split proportions must sum to 1, got {total}")
 
-            # Convert ratios to absolute sizes
+            # Calculate split sizes
             total_size = len(dataset)
-            if total_size == 0:
-                raise ValueError("Dataset is empty")
-                
-            train_end = int(total_size * train_size)
-            val_end = train_end + int(total_size * val_size)
+            train_samples = int(total_size * train_size)
+            val_samples = int(total_size * val_size)
+            test_samples = total_size - train_samples - val_samples
 
-            # Shuffle dataset with seed for reproducibility
-            dataset = dataset.shuffle(seed=42)
+            self.logger.log_info("Splitting dataset...")
 
-            # Split dataset
-            train_dataset = dataset.select(range(train_end))
-            val_dataset = dataset.select(range(train_end, val_end))
-            test_dataset = dataset.select(range(val_end, total_size))
+            # Create indices
+            indices = np.arange(total_size)
+            if shuffle:
+                with tqdm(total=1, desc="Shuffling dataset", unit="operation") as pbar:
+                    rng = np.random.default_rng(seed)
+                    rng.shuffle(indices)
+                    pbar.update(1)
 
-            # Validate split sizes
-            if len(train_dataset) == 0 or len(val_dataset) == 0 or len(test_dataset) == 0:
-                raise ValueError("One or more splits are empty. Try adjusting split ratios.")
+            # Split dataset using Hugging Face's built-in functionality
+            with tqdm(total=2, desc="Creating splits", unit="split") as pbar:
+                # First split: train vs rest
+                train_val_split = dataset.train_test_split(
+                    train_size=train_size,
+                    seed=seed,
+                    shuffle=False  # We already shuffled if needed
+                )
+                train_dataset = train_val_split["train"]
+                rest_dataset = train_val_split["test"]
+                pbar.update(1)
 
+                # Second split: val vs test from the rest
+                val_ratio = val_size / (val_size + test_size)
+                val_test_split = rest_dataset.train_test_split(
+                    train_size=val_ratio,
+                    seed=seed,
+                    shuffle=False
+                )
+                val_dataset = val_test_split["train"]
+                test_dataset = val_test_split["test"]
+                pbar.update(1)
+            
+            # Log split sizes
             self.logger.log_info(f"Split sizes - Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")
+            
             return train_dataset, val_dataset, test_dataset
 
         except Exception as e:
             self.logger.log_error(f"Error splitting dataset: {str(e)}")
             raise
-            
+
+    def train_val_split(
+        self,
+        dataset: Union[Dataset, DatasetDict],
+        train_size: float = 0.8,
+        shuffle: bool = True,
+        seed: int = 42,
+        split: str = "train"
+    ) -> Tuple[Dataset, Dataset]:
+        """Split dataset into train and validation sets."""
+        dataset = self._get_dataset_from_dict(dataset, split)
+        return dataset.train_test_split(
+            train_size=train_size,
+            shuffle=shuffle,
+            seed=seed
+        ).values()
+
     def k_fold_split(self, dataset, n_splits: int = 5, shuffle: bool = True, seed: int = 42):
         """Create k-fold cross validation splits."""
         try: