update client

modelscope · Oct 29, 2024 · 872f3f9 · 872f3f9
1 parent 7b4ad89
commit 872f3f9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/swift/llm/__init__.py b/swift/llm/__init__.py
@@ -6,7 +6,7 @@
 if TYPE_CHECKING:
     # Recommend using `xxx_main`
     from .infer import (VllmEngine, InferRequest, RequestConfig, InferStats, LmdeployEngine, PtEngine, infer_main,
-                        deploy_main, PtLoRARequest)
+                        deploy_main, PtLoRARequest, InferClient)
     from .export import export_main, merge_lora
     from .eval import eval_main
     from .train import sft_main, pt_main, rlhf_main
@@ -30,7 +30,7 @@
         'rlhf': ['rlhf_main'],
         'infer': [
             'deploy_main', 'VllmEngine', 'InferRequest', 'RequestConfig', 'InferStats', 'LmdeployEngine', 'PtEngine',
-            'infer_main', 'PtLoRARequest'
+            'infer_main', 'PtLoRARequest', 'InferClient'
         ],
         'export': ['export_main', 'merge_lora'],
         'eval': ['eval_main'],

diff --git a/swift/llm/infer/__init__.py b/swift/llm/infer/__init__.py
@@ -7,14 +7,16 @@
     from .infer import infer_main
     from .deploy import deploy_main
     from .protocol import InferRequest, RequestConfig
-    from .infer_engine import (InferEngine, VllmEngine, LmdeployEngine, PtEngine, InferStats, PtLoRARequest)
+    from .infer_engine import (InferEngine, VllmEngine, LmdeployEngine, PtEngine, InferStats, PtLoRARequest,
+                               InferClient)
 else:
     _extra_objects = {k: v for k, v in globals().items() if not k.startswith('_')}
     _import_structure = {
         'deploy': ['deploy_main'],
         'infer': ['infer_main'],
         'protocol': ['InferRequest', 'RequestConfig'],
-        'infer_engine': ['InferEngine', 'VllmEngine', 'LmdeployEngine', 'PtEngine', 'InferStats', 'PtLoRARequest'],
+        'infer_engine': ['InferEngine', 'VllmEngine', 'LmdeployEngine', 'PtEngine', 'InferStats', 'PtLoRARequest',
+                         'InferClient'],
     }
 
     import sys

diff --git a/swift/llm/infer/infer_engine/__init__.py b/swift/llm/infer/infer_engine/__init__.py
@@ -7,6 +7,7 @@
     from .vllm_engine import VllmEngine
     from .lmdeploy_engine import LmdeployEngine
     from .pt_engine import PtEngine, PtLoRARequest
+    from .infer_client import InferClient
     from .infer_engine import InferEngine
     from .base import BaseInferEngine
     from .utils import InferStats
@@ -16,9 +17,10 @@
         'vllm_engine': ['VllmEngine'],
         'lmdeploy_engine': ['LmdeployEngine'],
         'pt_engine': ['PtEngine', 'PtLoRARequest'],
+        'infer_client': ['InferClient'],
         'infer_engine': ['InferEngine'],
         'base': ['BaseInferEngine'],
-        'utils': ['InferStats']
+        'utils': ['InferStats'],
     }
 
     import sys