ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make the chunk function efficient

Browse source
This commit is contained in:
vangmay 2025-11-20 21:08:33 +08:00
commit a12eefa75a
3 changed files with 100 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

20
unsloth-cli.py
View file

@ -103,13 +103,19 @@ def run(args):
from transformers.utils import strtobool
if args.raw_text_file:
# Use raw text loader
# Use raw text loader - returns pre-tokenized data
loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride)
dataset = loader.load_from_file(args.raw_text_file)
dataset = loader.load_from_file(args.raw_text_file, return_tensors=True)
# Mark dataset as pre-tokenized to skip text formatting
dataset._is_pretokenized = True
return dataset
elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")):
# Auto-detect local raw text files
loader = RawTextDataLoader(tokenizer)
dataset = loader.load_from_file(args.dataset)
# Auto-detect local raw text files - returns pre-tokenized data
loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride)
dataset = loader.load_from_file(args.dataset, return_tensors=True)
# Mark dataset as pre-tokenized to skip text formatting
dataset._is_pretokenized = True
return dataset
else:
# Check for modelscope usage
use_modelscope = strtobool(
@ -123,9 +129,9 @@ def run(args):
# Existing HuggingFace dataset logic
dataset = load_dataset(args.dataset, split = "train")
# Apply formatting for structured datasets
# Apply formatting for structured datasets (text-based)
dataset = dataset.map(formatting_prompts_func, batched = True)
return dataset
return dataset
# Load dataset using smart loader
dataset = load_dataset_smart(args)