update

config.json

@@ -1,11 +1,11 @@
 {
     "architectures": [
-        "LmdeployForCausalLM"
+        "LMDeployForCausalLM"
     ],
     "auto_map": {
-        "AutoConfig": "configuration_lmdeploy.LmdeployConfig",
-        "AutoModel": "modeling_lmdeploy.LmdeployForCausalLM",
-        "AutoModelForCausalLM": "modeling_lmdeploy.LmdeployForCausalLM"
+        "AutoConfig": "configuration_lmdeploy.LMDeployConfig",
+        "AutoModel": "modeling_lmdeploy.LMDeployForCausalLM",
+        "AutoModelForCausalLM": "modeling_lmdeploy.LMDeployForCausalLM"
     },
     "turbomind": {
         "model_name": "internlm-chat-20b",
@@ -35,5 +35,6 @@
         "max_position_embeddings": 2048,
         "use_dynamic_ntk": 0,
         "use_logn_attn": 0
-    }
+    },
+    "lmdeploy_version": "0.0.14"
 }

configuration_lmdeploy.py

@@ -7,7 +7,8 @@ from lmdeploy.turbomind.deploy.target_model.base import TurbomindModelConfig
 from lmdeploy.version import __version__ as lm_version
 
 
-class LmdeployConfig(PretrainedConfig):
+class LMDeployConfig(PretrainedConfig):
+    """Lmdeploy config."""
 
     def __init__(self, turbomind: dict = None, **kwargs):
         default_tm_cfg = copy.deepcopy(
@@ -33,3 +34,4 @@ class LmdeployConfig(PretrainedConfig):
             return config, kwargs
         else:
             return config
+

modeling_lmdeploy.py

@@ -7,14 +7,15 @@ from itertools import count
 from queue import Queue
 from typing import List, Optional, Tuple, Union
 
+from huggingface_hub import snapshot_download
 from transformers import PretrainedConfig
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 
 from lmdeploy.turbomind import TurboMind
-from lmdeploy.turbomind.utils import download_hf_repo, get_gen_param
+from lmdeploy.turbomind.utils import get_gen_param
 
-from .configuration_lmdeploy import LmdeployConfig
+from .configuration_lmdeploy import LMDeployConfig
 
 logger = logging.get_logger(__name__)
 
@@ -55,11 +56,11 @@ class Session:
         return self._error
 
 
-class LmdeployForCausalLM(PreTrainedModel):
-    config_class = LmdeployConfig
+class LMDeployForCausalLM(PreTrainedModel):
+    config_class = LMDeployConfig
 
     def __init__(self,
-                 config: LmdeployConfig,
+                 config: LMDeployConfig,
                  *inputs,
                  model_path: str = None,
                  **kwargs):
@@ -90,7 +91,7 @@ class LmdeployForCausalLM(PreTrainedModel):
         if os.path.isdir(pretrained_model_name_or_path):
             local_folder = pretrained_model_name_or_path
         else:
-            local_folder = download_hf_repo(
+            local_folder = snapshot_download(
                 pretrained_model_name_or_path,
                 revision=revision,
                 cache_dir=cache_dir,
@@ -137,6 +138,7 @@ class LmdeployForCausalLM(PreTrainedModel):
                                             sequence_end=False,
                                             stop=True):
                 pass
+            session._error = 1
         finally:
             self.que.put(generator)
 
@@ -222,3 +224,4 @@ class LmdeployForCausalLM(PreTrainedModel):
                 session._step = _step + response_size
 
                 yield response, session
+