mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2025-12-30 02:30:35 +08:00
[v1] init data plugins (#9248)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Yaowei Zheng
GitHub
@@ -13,8 +13,8 @@
|
||||
# limitations under the License.
|
||||
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Literal, Optional
|
||||
from collections.abc import AsyncIterable, Iterable
|
||||
from typing import Any, Union
|
||||
|
||||
from datasets import load_dataset
|
||||
from huggingface_hub import hf_hub_download
|
||||
@@ -22,84 +22,22 @@ from omegaconf import OmegaConf
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from ..config.data_args import DataArguments
|
||||
from ..extras.types import DatasetInfo, HFDataset, Processor
|
||||
from ..extras.types import DatasetInfo, HFDataset, Processor, Tensor
|
||||
from ..plugins.data_plugins.loader import DataGetItemPlugin, DataIndexPlugin, DataLoaderPlugin
|
||||
|
||||
|
||||
class DataCollator:
|
||||
"""Default Data collator."""
|
||||
|
||||
def __init__(self, processor: Processor) -> None:
|
||||
def __init__(self, processor: Processor, dataset_info: dict[str, DatasetInfo]) -> None:
|
||||
self.processor = processor
|
||||
self.dataset_info = dataset_info
|
||||
|
||||
def __call__(self, features: list[dict[str, Any]]) -> dict[str, Tensor]:
|
||||
pass
|
||||
|
||||
|
||||
class DatasetPathMixin:
|
||||
"""Path utilities."""
|
||||
|
||||
args: DataArguments
|
||||
"""Data arguments."""
|
||||
|
||||
def _abspath(self, path: str, dataset_dir: Optional[str] = None) -> str:
|
||||
"""Get absolute path of dataset.
|
||||
|
||||
Args:
|
||||
path (str): Dataset path.
|
||||
dataset_dir (Optional[str], optional): Dataset directory. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: Absolute path of dataset.
|
||||
"""
|
||||
dataset_dir = dataset_dir or self.args.dataset_dir
|
||||
return os.path.abspath(os.path.expanduser(os.path.join(dataset_dir, path)))
|
||||
|
||||
def _exists(self, path: str, dataset_dir: Optional[str] = None) -> bool:
|
||||
"""Check if dataset exists.
|
||||
|
||||
Args:
|
||||
path (str): Dataset path.
|
||||
dataset_dir (Optional[str], optional): Dataset directory. Defaults to None.
|
||||
|
||||
Returns:
|
||||
bool: Whether dataset exists.
|
||||
"""
|
||||
return os.path.exists(self._abspath(path, dataset_dir))
|
||||
|
||||
def _isfile(self, path: str, dataset_dir: Optional[str] = None) -> bool:
|
||||
"""Check if dataset is a file.
|
||||
|
||||
Args:
|
||||
path (str): Dataset path.
|
||||
dataset_dir (Optional[str], optional): Dataset directory. Defaults to None.
|
||||
|
||||
Returns:
|
||||
bool: Whether dataset is a file.
|
||||
"""
|
||||
return os.path.isfile(self._abspath(path, dataset_dir))
|
||||
|
||||
def _isdir(self, path: str, dataset_dir: Optional[str] = None) -> bool:
|
||||
"""Check if dataset is a directory.
|
||||
|
||||
Args:
|
||||
path (str): Dataset path.
|
||||
dataset_dir (Optional[str], optional): Dataset directory. Defaults to None.
|
||||
|
||||
Returns:
|
||||
bool: Whether dataset is a directory.
|
||||
"""
|
||||
return os.path.isdir(self._abspath(path, dataset_dir))
|
||||
|
||||
def _get_builder_name(self, path: str) -> Literal["arrow", "csv", "json", "parquet", "text"]:
|
||||
"""Get dataset builder name.
|
||||
|
||||
Args:
|
||||
path (str): Dataset path.
|
||||
|
||||
Returns:
|
||||
Literal["arrow", "csv", "json", "parquet", "text"]: Dataset builder name.
|
||||
"""
|
||||
return os.path.splitext(path)[-1][1:].replace("jsonl", "json").replace("txt", "text")
|
||||
|
||||
|
||||
class DataEngine(Dataset, DatasetPathMixin):
|
||||
class DataEngine(Dataset):
|
||||
"""Data engine."""
|
||||
|
||||
def __init__(self, data_args: DataArguments) -> None:
|
||||
@@ -113,19 +51,27 @@ class DataEngine(Dataset, DatasetPathMixin):
|
||||
"""Whether dataset is streaming."""
|
||||
self.data_index: list[tuple[str, int]] = []
|
||||
"""List of (dataset_name, sample_index)"""
|
||||
self.data_loader_plugin = DataLoaderPlugin(args=self.args)
|
||||
"""Data loader plugin."""
|
||||
self.data_index_plugin = DataIndexPlugin()
|
||||
"""Data index plugin."""
|
||||
self.data_getitem_plugin = DataGetItemPlugin(datasets=self.datasets, data_index=self.data_index)
|
||||
"""Data getitem plugin."""
|
||||
self.get_dataset_info()
|
||||
self.load_dataset()
|
||||
self.build_data_index()
|
||||
|
||||
def get_dataset_info(self) -> None:
|
||||
"""Get dataset info."""
|
||||
if self.args.dataset.endswith(".yaml") and self._isfile(self.args.dataset): # local file
|
||||
self.dataset_info = OmegaConf.load(self._abspath(self.args.dataset))
|
||||
if self.args.dataset.endswith(".yaml") and os.path.isfile(
|
||||
os.path.join(self.args.dataset_dir, self.args.dataset)
|
||||
): # local file
|
||||
self.dataset_info = OmegaConf.load(os.path.join(self.args.dataset_dir, self.args.dataset))
|
||||
elif self.args.dataset.endswith(".yaml"): # hf hub uri, e.g. llamafactory/v1-sft-demo/dataset_info.yaml
|
||||
repo_id, filename = os.path.split(self.args.dataset)
|
||||
filepath = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset")
|
||||
self.dataset_info = OmegaConf.load(filepath)
|
||||
elif self._exists(self.args.dataset): # local file(s)
|
||||
elif os.path.exists(os.path.join(self.args.dataset_dir, self.args.dataset)): # local file(s)
|
||||
self.dataset_info = {"default": {"file_name": self.args.dataset}}
|
||||
else: # hf hub dataset, e.g. llamafactory/v1-sft-demo
|
||||
self.dataset_info = {"default": {"hf_hub_url": self.args.dataset}}
|
||||
@@ -133,37 +79,39 @@ class DataEngine(Dataset, DatasetPathMixin):
|
||||
def load_dataset(self) -> None:
|
||||
"""Load dataset from dataset info."""
|
||||
for key, value in self.dataset_info.items():
|
||||
dataset_dir = value.get("dataset_dir", self.args.dataset_dir)
|
||||
split = value.get("split", "train")
|
||||
streaming = value.get("streaming", False)
|
||||
self.streaming |= streaming
|
||||
if "hf_hub_url" in value:
|
||||
self.datasets[key] = load_dataset(value["hf_hub_url"], split=split, streaming=streaming)
|
||||
elif "file_name" in value:
|
||||
filepath = self._abspath(value["file_name"], dataset_dir)
|
||||
if os.path.isdir(filepath):
|
||||
filetype = self._get_builder_name(os.listdir(filepath)[0])
|
||||
self.datasets[key] = load_dataset(filetype, data_dir=filepath, split=split)
|
||||
elif os.path.isfile(filepath):
|
||||
filetype = self._get_builder_name(filepath)
|
||||
self.datasets[key] = load_dataset(filetype, data_files=filepath, split=split)
|
||||
else:
|
||||
raise ValueError(f"Can not load dataset {key} from {filepath}.")
|
||||
|
||||
if streaming:
|
||||
self.datasets[key] = self.datasets[key].to_iterable_dataset()
|
||||
else:
|
||||
# TODO: support dataset loader plugins
|
||||
raise ValueError(f"Dataset {key} is not supported.")
|
||||
else: # data loader plugin
|
||||
self.datasets[key] = self.data_loader_plugin.auto_load_data(value)
|
||||
|
||||
def build_data_index(self) -> None:
|
||||
"""Build dataset index."""
|
||||
for dataset_name, dataset in self.datasets.items():
|
||||
size = self.dataset_info[dataset_name].get("size")
|
||||
weight = self.dataset_info[dataset_name].get("weight")
|
||||
if self.streaming:
|
||||
self.data_index.append((dataset_name, -1))
|
||||
data_index = [(dataset_name, -1) for _ in range(1000)]
|
||||
else:
|
||||
# TODO: add sample_num, weight
|
||||
self.data_index.extend([(dataset_name, sample_index) for sample_index in range(len(dataset))])
|
||||
data_index = [(dataset_name, sample_index) for sample_index in range(len(dataset))]
|
||||
|
||||
if size or weight: # data index plugin
|
||||
data_index = self.data_index_plugin.adjust_data_index(data_index, size, weight)
|
||||
|
||||
self.data_index.extend(data_index)
|
||||
|
||||
def get_data_collator(self, processor: Processor) -> DataCollator:
|
||||
"""Get data collator.
|
||||
|
||||
Args:
|
||||
processor (Processor): Processor.
|
||||
|
||||
Returns:
|
||||
DataCollator: Data collator.
|
||||
"""
|
||||
return DataCollator(processor=processor, dataset_info=self.dataset_info)
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Get dataset length.
|
||||
@@ -176,7 +124,7 @@ class DataEngine(Dataset, DatasetPathMixin):
|
||||
else:
|
||||
return len(self.data_index)
|
||||
|
||||
def __getitem__(self, index: int) -> dict:
|
||||
def __getitem__(self, index: Union[int, slice, list[int]]) -> Union[dict, list[dict]]:
|
||||
"""Get dataset item.
|
||||
|
||||
Args:
|
||||
@@ -185,21 +133,39 @@ class DataEngine(Dataset, DatasetPathMixin):
|
||||
Returns:
|
||||
dict: Dataset item.
|
||||
"""
|
||||
dataset_name, sample_index = self.data_index[index]
|
||||
return self.datasets[dataset_name][sample_index]
|
||||
if self.streaming:
|
||||
raise ValueError("Streaming dataset does not support index access.")
|
||||
|
||||
def __iter__(self) -> Iterator:
|
||||
if isinstance(index, int):
|
||||
dataset_name, sample_index = self.data_index[index]
|
||||
return {"_dataset_name": dataset_name, **self.datasets[dataset_name][sample_index]}
|
||||
else:
|
||||
return self.data_getitem_plugin.get_data(index)
|
||||
|
||||
def __iter__(self) -> Iterable:
|
||||
"""Get dataset iterator.
|
||||
|
||||
Returns:
|
||||
Iterator: Dataset iterator.
|
||||
Iterable: Dataset iterator.
|
||||
"""
|
||||
if self.streaming:
|
||||
pass
|
||||
else:
|
||||
# TODO: add shuffle here
|
||||
pass
|
||||
|
||||
raise NotImplementedError()
|
||||
|
||||
def __aiter__(self) -> AsyncIterator:
|
||||
def __aiter__(self) -> AsyncIterable:
|
||||
"""Get dataset async iterator.
|
||||
|
||||
Returns:
|
||||
AsyncIterator: Dataset async iterator.
|
||||
AsyncIterable: Dataset async iterator.
|
||||
"""
|
||||
if self.streaming:
|
||||
pass
|
||||
else:
|
||||
# TODO: add shuffle here
|
||||
pass
|
||||
|
||||
raise NotImplementedError()
|
||||
|
||||
Reference in New Issue
Block a user