加入MOSS支持

modules/models/MOSS.py

@@ -0,0 +1,33 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+import time
+import numpy as np
+from torch.nn import functional as F
+import os
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from transformers import MossForCausalLM, MossConfig
+
+from .base_model import BaseLLMModel
+
+MOSS_MODEL = None
+MOSS_TOKENIZER = None
+
+class MOSS_Client(BaseLLMModel):
+    def __init__(self, model_name) -> None:
+        super().__init__(model_name=model_name)
+        global MOSS_MODEL, MOSS_TOKENIZER
+        config = MossConfig.from_pretrained("fnlp/moss-16B-sft")
+        print("MOSS Model Parallelism Devices: ", torch.cuda.device_count())
+        with init_empty_weights():
+            raw_model = MossForCausalLM._from_config(config, torch_dtype=torch.float16)
+        raw_model.tie_weights()
+        MOSS_MODEL = load_checkpoint_and_dispatch(
+            raw_model,
+            "fnlp/moss-16B-sft",
+            device_map="auto",
+            no_split_module_classes=["MossBlock"],
+            dtype=torch.float16
+        )
+
+if __name__ == "__main__":
+    model = MOSS_Client("MOSS")

modules/models/base_model.py

@@ -31,6 +31,7 @@ class ModelType(Enum):
     LLaMA = 2
     XMChat = 3
     StableLM = 4
+    MOSS = 5
 
     @classmethod
     def get_type(cls, model_name: str):
@@ -46,6 +47,8 @@ class ModelType(Enum):
             model_type = ModelType.XMChat
         elif "stablelm" in model_name_lower:
             model_type = ModelType.StableLM
+        elif "moss" in model_name_lower:
+            model_type = ModelType.MOSS
         else:
             model_type = ModelType.Unknown
         return model_type

modules/presets.py

@@ -74,11 +74,12 @@ LOCAL_MODELS = [
     "chatglm-6b",
     "chatglm-6b-int4",
     "chatglm-6b-int4-qe",
+    "StableLM",
+    "MOSS",
     "llama-7b-hf",
     "llama-13b-hf",
     "llama-30b-hf",
     "llama-65b-hf",
-    "StableLM"
 ]
 
 if os.environ.get('HIDE_LOCAL_MODELS', 'false') == 'true':