From f80fe573e6a44e1337cfd2bae7022fba53437d44 Mon Sep 17 00:00:00 2001 From: vangmay Date: Thu, 20 Nov 2025 20:53:22 +0800 Subject: [PATCH] Integrate smart dataset loader --- unsloth-cli.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/unsloth-cli.py b/unsloth-cli.py index aac0e7f7e1..044ad93f31 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -100,6 +100,8 @@ def run(args): return {"text": texts} def load_dataset_smart(args): + from transformers.utils import strtobool + if args.raw_text_file: # Use raw text loader loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride) @@ -109,20 +111,21 @@ def run(args): loader = RawTextDataLoader(tokenizer) dataset = loader.load_from_file(args.dataset) else: - # Existing HuggingFace dataset logic - dataset = load_dataset(args.dataset, split="train") + # Check for modelscope usage + use_modelscope = strtobool(os.environ.get("UNSLOTH_USE_MODELSCOPE", "False")) + if use_modelscope: + from modelscope import MsDataset + dataset = MsDataset.load(args.dataset, split="train") + else: + # Existing HuggingFace dataset logic + dataset = load_dataset(args.dataset, split="train") + + # Apply formatting for structured datasets dataset = dataset.map(formatting_prompts_func, batched=True) return dataset - use_modelscope = strtobool(os.environ.get("UNSLOTH_USE_MODELSCOPE", "False")) - if use_modelscope: - from modelscope import MsDataset - - dataset = MsDataset.load(args.dataset, split = "train") - else: - # Load and format dataset - dataset = load_dataset(args.dataset, split = "train") - dataset = dataset.map(formatting_prompts_func, batched = True) + # Load dataset using smart loader + dataset = load_dataset_smart(args) print("Data is formatted and ready!") # Configure training arguments