Update arguments (#162)

Author: Jintao-Huang

swift/llm/utils/argument.py

@@ -517,6 +517,8 @@ def register_custom_dataset(args: Union[SftArguments, InferArguments]) -> None:
     if args.custom_train_dataset_path is None:
         assert args.custom_val_dataset_path is None
         return
+    if '_custom_dataset' in DATASET_MAPPING:
+        DATASET_MAPPING.pop('_custom_dataset')
     register_dataset(
         '_custom_dataset',
         '_custom_dataset',
@@ -549,7 +551,10 @@ def load_from_ckpt_dir(args: InferArguments) -> None:
             'custom_val_dataset_path'
         ]
     for key in imported_keys:
-        if key == 'model_cache_dir' and getattr(args, key) is not None:
+        if (key in {
+                'model_cache_dir', 'dataset', 'custom_train_dataset_path',
+                'custom_val_dataset_path'
+        } and getattr(args, key) is not None):
             continue
         setattr(args, key, sft_args.get(key))