Support swift in llm_pipeline. (#595)

* support swift in llm_pipeline

* adjust function name

modelscope/pipeline_inputs.py

@@ -326,6 +326,20 @@ TASK_INPUTS = {
 
     # ============ nlp tasks ===================
     Tasks.chat: {
+        # An input example for `messages` format (Dict[str, List[Dict[str, str]]]):
+        # {'messages': [{
+        #     'role': 'system',
+        #     'content': 'You are a helpful assistant.'
+        # }, {
+        #     'role': 'user',
+        #     'content': 'Hello! Where is the capital of Zhejiang?'
+        # }, {
+        #     'role': 'assistant',
+        #     'content': 'Hangzhou is the capital of Zhejiang.'
+        # }, {
+        #     'role': 'user',
+        #     'content': 'Tell me something about HangZhou?'
+        # }]}
         'messages': InputType.LIST
     },
     Tasks.text_classification: [

modelscope/pipelines/builder.py

@@ -215,7 +215,7 @@ def llm_first_checker(model: Union[str, List[str], Model, List[Model]],
         except Exception:
             return None
 
-    def parse_model_type(file: Optional[str], pattern: str) -> Optional[str]:
+    def parse_and_get(file: Optional[str], pattern: str) -> Optional[str]:
         if file is None or not osp.exists(file):
             return None
         return Config.from_file(file).safe_get(pattern)
@@ -223,15 +223,22 @@ def llm_first_checker(model: Union[str, List[str], Model, List[Model]],
     def get_model_type(model: str, revision: Optional[str]) -> Optional[str]:
         cfg_file = get_file_name(model, ModelFile.CONFIGURATION, revision)
         hf_cfg_file = get_file_name(model, ModelFile.CONFIG, revision)
-        cfg_model_type = parse_model_type(cfg_file, 'model.type')
-        hf_cfg_model_type = parse_model_type(hf_cfg_file, 'model_type')
+        cfg_model_type = parse_and_get(cfg_file, 'model.type')
+        hf_cfg_model_type = parse_and_get(hf_cfg_file, 'model_type')
         return cfg_model_type or hf_cfg_model_type
 
+    def get_adapter_type(model: str, revision: Optional[str]) -> Optional[str]:
+        cfg_file = get_file_name(model, ModelFile.CONFIGURATION, revision)
+        model = parse_and_get(cfg_file, 'adapter_cfg.model_id_or_path')
+        revision = parse_and_get(cfg_file, 'adapter_cfg.model_revision')
+        return None if model is None else get_model_type(model, revision)
+
     if isinstance(model, list):
         model = model[0]
     if not isinstance(model, str):
         model = model.model_dir
-    model_type = get_model_type(model, revision)
+    model_type = get_model_type(model, revision) \
+        or get_adapter_type(model, revision)
     if model_type is not None:
         model_type = model_type.lower().split('-')[0]
         if model_type in LLM_FORMAT_MAP:

modelscope/pipelines/nlp/llm_pipeline.py

@@ -9,11 +9,13 @@ from transformers import PreTrainedTokenizer
 
 from modelscope import (AutoModelForCausalLM, AutoTokenizer, Pipeline,
                         snapshot_download)
+from modelscope.hub.file_download import model_file_download
 from modelscope.models.base import Model
 from modelscope.models.nlp import ChatGLM2Tokenizer, Llama2Tokenizer
 from modelscope.outputs import OutputKeys
 from modelscope.pipelines.builder import PIPELINES
 from modelscope.pipelines.util import is_model, is_official_hub_path
+from modelscope.utils.config import Config
 from modelscope.utils.constant import Invoke, ModelFile, Tasks
 from modelscope.utils.logger import get_logger
 
@@ -27,6 +29,22 @@ class LLMPipeline(Pipeline):
     def initiate_single_model(self, model):
         if isinstance(model, str):
             logger.info(f'initiate model from {model}')
+        if self._is_swift_model(model):
+            from swift import Swift
+
+            base_model = self.cfg.safe_get('adapter_cfg.model_id_or_path')
+            assert base_model is not None, 'Cannot get adapter_cfg.model_id_or_path from configuration.json file.'
+            revision = self.cfg.safe_get('adapter_cfg.model_revision',
+                                         'master')
+            base_model = Model.from_pretrained(
+                base_model,
+                revision,
+                invoked_by=Invoke.PIPELINE,
+                device_map=self.device_map,
+                torch_dtype=self.torch_dtype,
+                trust_remote_code=True)
+            swift_model = Swift.from_pretrained(base_model, model_id=model)
+            return swift_model
         if isinstance(model, str) and is_official_hub_path(model):
             logger.info(f'initiate model from location {model}.')
             if is_model(model):
@@ -50,6 +68,20 @@ class LLMPipeline(Pipeline):
         else:
             return model
 
+    def _is_swift_model(self, model: Union[str, Any]) -> bool:
+        if not isinstance(model, str):
+            return False
+        if os.path.exists(model):
+            cfg_file = os.path.join(model, ModelFile.CONFIGURATION)
+        else:
+            try:
+                cfg_file = model_file_download(model, ModelFile.CONFIGURATION)
+            except Exception:
+                return False
+
+        self.cfg = Config.from_file(cfg_file)
+        return self.cfg.safe_get('adapter_cfg.tuner_backend') == 'swift'
+
     def __init__(self,
                  format_messages: Union[Callable, str] = None,
                  format_output: Callable = None,

tests/pipelines/test_base.py

@@ -166,12 +166,24 @@ class CustomPipelineTest(unittest.TestCase):
                 return inputs
 
             def postprocess(self, out, **kwargs):
-                return {'response': 'xxx', 'history': []}
+                return {'message': {'role': 'assistant', 'content': 'xxx'}}
 
         pipe = pipeline(
             task=Tasks.chat, pipeline_name=dummy_module, model=self.model_dir)
         pipe('text')
         inputs = {'text': 'aaa', 'history': [('dfd', 'fds')]}
+        inputs = {
+            'messages': [{
+                'role': 'user',
+                'content': 'dfd'
+            }, {
+                'role': 'assistant',
+                'content': 'fds'
+            }, {
+                'role': 'user',
+                'content': 'aaa'
+            }]
+        }
         pipe(inputs)
 
     def test_custom(self):