add result_dir paramerter to InferArgument (#1561)

Author: starxhong

swift/llm/infer.py

@@ -325,11 +325,15 @@ def llm_infer(args: InferArguments) -> Dict[str, List[Dict[str, Any]]]:
     result: List[Dict[str, Any]] = []
     jsonl_path = None
     if args.save_result:
-        result_dir = args.ckpt_dir
-        if result_dir is None:
-            result_dir = llm_engine.model_dir if args.infer_backend in {'vllm', 'lmdeploy'} else model.model_dir
+        if args.result_dir:
+            result_dir = args.result_dir
+        else:
+            result_dir = args.ckpt_dir
+            if result_dir is None:
+                result_dir = llm_engine.model_dir if args.infer_backend in {'vllm', 'lmdeploy'} else model.model_dir
+            if result_dir is not None:
+                result_dir = os.path.join(result_dir, 'infer_result')
         if result_dir is not None:
-            result_dir = os.path.join(result_dir, 'infer_result')
             os.makedirs(result_dir, exist_ok=True)
             time = dt.datetime.now().strftime('%Y%m%d-%H%M%S')
             jsonl_path = os.path.join(result_dir, f'{time}.jsonl')

swift/llm/utils/argument.py

@@ -1163,6 +1163,7 @@ class InferArguments(ArgumentsBase):
         default='AUTO', metadata={'help': f"template_type choices: {list(TEMPLATE_MAPPING.keys()) + ['AUTO']}"})
     infer_backend: Literal['AUTO', 'vllm', 'pt', 'lmdeploy'] = 'AUTO'
     ckpt_dir: Optional[str] = field(default=None, metadata={'help': '/path/to/your/vx-xxx/checkpoint-xxx'})
+    result_dir: Optional[str] = field(default=None, metadata={'help': '/path/to/your/infer_result'})
     load_args_from_ckpt_dir: bool = True
     load_dataset_config: bool = False
     eval_human: Optional[bool] = None

swift/llm/utils/utils.py

@@ -613,7 +613,7 @@ def _prepare_inputs(model: PreTrainedModel,
         if max_length and token_len + generation_config.max_new_tokens > max_length:
             generation_config.max_new_tokens = max_length - token_len
             if generation_config.max_new_tokens <= 0:
-                raise AssertionError('Current sentence length exceeds' f'the model max_length: {max_length}')
+                raise AssertionError(f'Current sentence length exceeds the model max_length: {max_length}')
     if template.suffix[-1] not in stop_words:
         stop_words.append(template.suffix[-1])
     inputs = to_device(inputs, device)