""" SmolLM3 Dataset Handler Handles data loading, preprocessing, and tokenization for SmolLM3 fine-tuning """ import os import json import torch from typing import Dict, List, Optional, Union from datasets import Dataset, load_dataset from transformers import PreTrainedTokenizer import logging logger = logging.getLogger(__name__) class SmolLM3Dataset: """Dataset handler for SmolLM3 fine-tuning""" def __init__( self, data_path: str, tokenizer: PreTrainedTokenizer, max_seq_length: int = 4096, use_chat_template: bool = True, chat_template_kwargs: Optional[Dict] = None, filter_bad_entries: bool = False, bad_entry_field: str = "bad_entry", sample_size: Optional[int] = None, sample_seed: int = 42 ): self.data_path = data_path self.tokenizer = tokenizer self.max_seq_length = max_seq_length self.use_chat_template = use_chat_template self.chat_template_kwargs = chat_template_kwargs or {} self.filter_bad_entries = filter_bad_entries self.bad_entry_field = bad_entry_field self.sample_size = sample_size self.sample_seed = sample_seed # Load and process dataset self.dataset = self._load_dataset() self.processed_dataset = self._process_dataset() def _load_dataset(self) -> Dataset: """Load dataset from various formats""" logger.info("Loading dataset from %s", self.data_path) # Check if it's a Hugging Face dataset if os.path.isdir(self.data_path): # Local directory try: dataset = load_dataset("json", data_files={ "train": os.path.join(self.data_path, "train.json"), "validation": os.path.join(self.data_path, "validation.json") if os.path.exists(os.path.join(self.data_path, "validation.json")) else None, "test": os.path.join(self.data_path, "test.json") if os.path.exists(os.path.join(self.data_path, "test.json")) else None }) logger.info("Loaded dataset from local JSON files") return dataset except Exception as e: logger.warning("Failed to load as JSON dataset: %s", e) # Try to load as a single JSON file if os.path.isfile(self.data_path) and self.data_path.endswith('.json'): try: with open(self.data_path, 'r', encoding='utf-8') as f: data = json.load(f) # Convert to dataset format if isinstance(data, list): dataset = Dataset.from_list(data) else: dataset = Dataset.from_dict(data) logger.info("Loaded dataset from single JSON file") return dataset except Exception as e: logger.error("Failed to load JSON file: %s", e) raise # Try to load as a Hugging Face dataset name try: dataset = load_dataset(self.data_path) logger.info("Loaded Hugging Face dataset: %s", self.data_path) # Filter bad entries if requested if self.filter_bad_entries and self.bad_entry_field in dataset["train"].column_names: logger.info("Filtering out bad entries using field: %s", self.bad_entry_field) for split in dataset: if self.bad_entry_field in dataset[split].column_names: original_size = len(dataset[split]) dataset[split] = dataset[split].filter(lambda x: not x[self.bad_entry_field]) filtered_size = len(dataset[split]) logger.info("Filtered %s: %d -> %d samples", split, original_size, filtered_size) # Apply sampling if requested if self.sample_size is not None and "train" in dataset: logger.info(f"Sampling {self.sample_size} random samples from {len(dataset['train'])} total samples") import random random.seed(self.sample_seed) # Sample indices total_samples = len(dataset["train"]) if self.sample_size > total_samples: logger.warning(f"Requested sample size ({self.sample_size}) is larger than dataset size ({total_samples}). Using all samples.") sampled_indices = list(range(total_samples)) else: sampled_indices = random.sample(range(total_samples), self.sample_size) # Apply sampling to train split dataset["train"] = dataset["train"].select(sampled_indices) logger.info(f"Sampled {len(dataset['train'])} train samples") # Also sample validation if it exists and is large if "validation" in dataset and len(dataset["validation"]) > 1000: val_sample_size = min(1000, len(dataset["validation"])) logger.info(f"Sampling {val_sample_size} validation samples from {len(dataset['validation'])} total") val_sampled_indices = random.sample(range(len(dataset["validation"])), val_sample_size) dataset["validation"] = dataset["validation"].select(val_sampled_indices) logger.info(f"Sampled {len(dataset['validation'])} validation samples") # If only 'train' split exists, create validation and test splits if ("train" in dataset) and ("validation" not in dataset or "test" not in dataset): logger.info("Automatically splitting train into train/validation/test (98/1/1)") split_dataset = dataset["train"].train_test_split(test_size=0.02, seed=42) # Now split test into validation and test (1% each) val_test_split = split_dataset["test"].train_test_split(test_size=0.5, seed=42) dataset = { "train": split_dataset["train"], "validation": val_test_split["train"], "test": val_test_split["test"] } return dataset except Exception as e: logger.error("Failed to load dataset: %s", e) raise def _process_dataset(self) -> Dataset: """Process the dataset for training""" logger.info("Processing dataset for training") def format_chat_template(example): """Format example using chat template""" if self.use_chat_template: try: # Handle different input formats if "messages" in example: messages = example["messages"] elif "conversations" in example: messages = example["conversations"] elif "user" in example and "assistant" in example: messages = [ {"role": "user", "content": example["user"]}, {"role": "assistant", "content": example["assistant"]} ] elif "instruction" in example and "output" in example: messages = [ {"role": "user", "content": example["instruction"]}, {"role": "assistant", "content": example["output"]} ] elif "prompt" in example and "completion" in example: messages = [ {"role": "user", "content": example["prompt"]}, {"role": "assistant", "content": example["completion"]} ] elif "prompt" in example and "accepted_completion" in example: messages = [ {"role": "user", "content": example["prompt"]}, {"role": "assistant", "content": example["accepted_completion"]} ] elif "prompt" in example and "completion" in example: messages = [ {"role": "user", "content": example["prompt"]}, {"role": "assistant", "content": example["completion"]} ] else: # Fallback: treat as plain text return {"text": str(example)} # Add system message with /no_think tag if not present if messages and messages[0]["role"] != "system": # Check if we should add /no_think tag based on configuration system_content = "Tu es TonicIA, un assistant francophone rigoureux et bienveillant." if hasattr(self, 'chat_template_kwargs') and self.chat_template_kwargs: # If no_think_system_message is True, add /no_think tag if self.chat_template_kwargs.get("no_think_system_message") == True: system_content = "Tu es TonicIA , un assistant francophone rigoureux et bienveillant. /no_think" messages.insert(0, {"role": "system", "content": system_content}) # Apply chat template text = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=self.chat_template_kwargs.get("add_generation_prompt", True) ) return {"text": text} except Exception as e: logger.warning("Failed to apply chat template: %s", e) # Fallback to plain text return {"text": str(example)} else: # Use plain text if "text" in example: return {"text": example["text"]} else: return {"text": str(example)} def tokenize_function(examples): """Tokenize the examples""" # Tokenize the texts with fixed length tokenized = self.tokenizer( examples["text"], truncation=True, padding=True, # Enable padding during tokenization max_length=self.max_seq_length, return_overflowing_tokens=False, # Don't return overflowing tokens return_length=True, ) # Calculate input length input_length = [len(x) for x in tokenized["input_ids"]] # Create labels (same as input_ids for causal LM) tokenized["labels"] = tokenized["input_ids"].copy() return { "input_ids": tokenized["input_ids"], "attention_mask": tokenized["attention_mask"], "labels": tokenized["labels"], "length": input_length, } # Process the dataset - handle both single dataset and dictionary of splits if isinstance(self.dataset, dict): # Process each split individually processed_dataset = {} for split_name, split_dataset in self.dataset.items(): logger.info("Processing %s split...", split_name) # Format the split processed_split = split_dataset.map( format_chat_template, remove_columns=split_dataset.column_names, desc="Formatting {} dataset".format(split_name) ) # Tokenize the split tokenized_split = processed_split.map( tokenize_function, remove_columns=processed_split.column_names, desc="Tokenizing {} dataset".format(split_name), batched=True, ) processed_dataset[split_name] = tokenized_split else: # Single dataset processed_dataset = self.dataset.map( format_chat_template, remove_columns=self.dataset.column_names, desc="Formatting dataset" ) # Tokenize the dataset processed_dataset = processed_dataset.map( tokenize_function, remove_columns=processed_dataset.column_names, desc="Tokenizing dataset", batched=True, ) # Log processing results if isinstance(processed_dataset, dict): logger.info("Dataset processed. Train samples: %d", len(processed_dataset['train'])) if "validation" in processed_dataset: logger.info("Validation samples: %d", len(processed_dataset['validation'])) if "test" in processed_dataset: logger.info("Test samples: %d", len(processed_dataset['test'])) else: logger.info("Dataset processed. Samples: %d", len(processed_dataset)) return processed_dataset def get_train_dataset(self) -> Dataset: """Get training dataset""" return self.processed_dataset["train"] def get_eval_dataset(self) -> Optional[Dataset]: """Get evaluation dataset if available""" if "validation" in self.processed_dataset: return self.processed_dataset["validation"] elif "test" in self.processed_dataset: return self.processed_dataset["test"] else: return None def get_data_collator(self): """Get data collator for training""" from transformers import DataCollatorForLanguageModeling import torch base_collator = DataCollatorForLanguageModeling( tokenizer=self.tokenizer, mlm=False, pad_to_multiple_of=8, return_tensors="pt", ) def collator_with_stats(features): batch = base_collator(features) # Calculate token stats input_ids = batch["input_ids"] attention_mask = batch.get("attention_mask", None) labels = batch.get("labels", None) pad_token_id = self.tokenizer.pad_token_id if pad_token_id is None: pad_token_id = self.tokenizer.eos_token_id total_tokens = int((input_ids != pad_token_id).sum().item()) padding_tokens = int((input_ids == pad_token_id).sum().item()) batch_size, seq_len = input_ids.shape # Truncated tokens: count tokens that were cut off due to max_seq_length # (Assume all input is truncated to max_seq_length, so count tokens at max length) truncated_tokens = 0 for f in features: if "length" in f and f["length"] >= self.max_seq_length: truncated_tokens += f["length"] - self.max_seq_length + 1 batch["total_tokens"] = total_tokens batch["padding_tokens"] = padding_tokens batch["truncated_tokens"] = truncated_tokens batch["batch_size"] = batch_size batch["seq_len"] = seq_len return batch return collator_with_stats def create_sample_dataset(output_path: str = "my_dataset"): """Create a sample dataset for testing""" os.makedirs(output_path, exist_ok=True) # Sample conversations conversations = [ { "messages": [ {"role": "user", "content": "What is machine learning?"}, {"role": "assistant", "content": "Machine learning is a subset of artificial intelligence that enables computers to learn and improve from experience without being explicitly programmed."} ] }, { "messages": [ {"role": "user", "content": "Explain gravity in simple terms."}, {"role": "assistant", "content": "Gravity is the force that pulls objects toward each other, like how the Earth pulls things down to the ground."} ] }, { "messages": [ {"role": "user", "content": "How do I make a cup of coffee?"}, {"role": "assistant", "content": "To make a cup of coffee: 1) Boil water, 2) Add coffee grounds to a filter, 3) Pour hot water over the grounds, 4) Let it brew for a few minutes, 5) Enjoy!"} ] } ] # Split into train/validation train_data = conversations[:2] validation_data = conversations[2:] # Save to files with open(os.path.join(output_path, "train.json"), 'w', encoding='utf-8') as f: json.dump(train_data, f, indent=2, ensure_ascii=False) with open(os.path.join(output_path, "validation.json"), 'w', encoding='utf-8') as f: json.dump(validation_data, f, indent=2, ensure_ascii=False) logger.info("Sample dataset created in %s", output_path) return output_path