423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 19:30:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

optimize data loading logic

Browse Source
This commit is contained in:
hiyouga
2023-12-20 16:15:41 +08:00
parent c6abbbfe90
commit ec1fe1daa9
2 changed files with 9 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/llmtuner/data/loader.py
View File

@@ -1,7 +1,7 @@
import os
from typing import TYPE_CHECKING, Any, Dict, List, Union
from datasets import concatenate_datasets, interleave_datasets, load_dataset
from datasets import concatenate_datasets, interleave_datasets, load_dataset, load_from_disk
from llmtuner.data.utils import checksum
from llmtuner.extras.constants import FILEEXT2TYPE
@@ -22,6 +22,13 @@ def get_dataset(
max_samples = data_args.max_samples
all_datasets: List[Union["Dataset", "IterableDataset"]] = [] # support multiple datasets
if data_args.cache_path is not None and os.path.exists(data_args.cache_path):
logger.warning("Loading dataset from disk will ignore other data arguments.")
dataset = load_from_disk(data_args.cache_path)
if data_args.streaming:
dataset = dataset.to_iterable_dataset()
return dataset
for dataset_attr in data_args.dataset_list:
logger.info("Loading dataset {}...".format(dataset_attr))