在UI中加入StableLM

modules/models/base_model.py

@@ -30,6 +30,7 @@ class ModelType(Enum):
     ChatGLM = 1
     LLaMA = 2
     XMChat = 3
+    StableLM = 4
 
     @classmethod
     def get_type(cls, model_name: str):
@@ -43,6 +44,8 @@ class ModelType(Enum):
             model_type = ModelType.LLaMA
         elif "xmchat" in model_name_lower:
             model_type = ModelType.XMChat
+        elif "StableLM" in model_name_lower:
+            model_type = ModelType.StableLM
         else:
             model_type = ModelType.Unknown
         return model_type

modules/models/models.py

@@ -577,6 +577,9 @@ def get_model(
             if os.environ.get("XMCHAT_API_KEY") != "":
                 access_key = os.environ.get("XMCHAT_API_KEY")
             model = XMChat(api_key=access_key)
+        elif model_type == ModelType.StableLM:
+            from StableLM import StableLM_Client
+            model = StableLM_Client(model_name)
         elif model_type == ModelType.Unknown:
             raise ValueError(f"未知模型: {model_name}")
         logging.info(msg)

modules/presets.py

@@ -77,7 +77,8 @@ LOCAL_MODELS = [
     "llama-7b-hf",
     "llama-13b-hf",
     "llama-30b-hf",
-    "llama-65b-hf"
+    "llama-65b-hf",
+    "StableLM"
 ]
 
 if os.environ.get('HIDE_LOCAL_MODELS', 'false') == 'true':