support layerwise galore

Former-commit-id: d43a4da0947897d0be3f62fad3107754d4c89f2b
2024-03-10 00:24:11 +08:00
parent c635bbe465
commit 7ff8a064f3
14 changed files with 109 additions and 51 deletions
--- a/src/llmtuner/data/loader.py
+++ b/src/llmtuner/data/loader.py
@@ -29,7 +29,7 @@ def load_single_dataset(
    dataset_attr: "DatasetAttr",
    model_args: "ModelArguments",
    data_args: "DataArguments",
-):
+) -> Union["Dataset", "IterableDataset"]:
    logger.info("Loading dataset {}...".format(dataset_attr))
    data_path, data_name, data_dir, data_files = None, None, None, None
    if dataset_attr.load_from in ["hf_hub", "ms_hub"]: