Merge pull request #2 from OpenLMLab/main

Release

Author: 00INDEX

.github/workflows/main.yml

@@ -55,7 +55,7 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         with:
           tag_name: ui-v${{ steps.version.outputs.value }}
-          release_name: Release refs/heads/ui
+          release_name: Release refs/heads/ui-v${{ steps.version.outputs.value }}
           body: UI version ${{ steps.version.outputs.value }}.
           draft: false
           prerelease: false

README.md

@@ -62,6 +62,9 @@ python server.py --pretrained_path fnlp/moss-moon-003-sft
 - [GODEL](https://github.com/microsoft/GODEL)
     - [GODEL-v1_1-base-seq2seq](https://huggingface.co/microsoft/GODEL-v1_1-base-seq2seq)
     - [GODEL-v1_1-large-seq2seq](https://huggingface.co/microsoft/GODEL-v1_1-large-seq2seq)
+- [StableLM]
+    - [stablelm-tuned-alpha-3b](https://huggingface.co/stabilityai/stablelm-tuned-alpha-3b)
+    - [stablelm-tuned-alpha-7b](https://huggingface.co/stabilityai/stablelm-tuned-alpha-7b)
 
 ### 添加自己的模型
 

config.py

@@ -29,6 +29,9 @@
     "microsoft/GODEL-v1_1-large-seq2seq": "godel",
     # belle
     "BelleGroup/BELLE-7B-2M": "belle",
+    # stablelm
+    "stabilityai/stablelm-tuned-alpha-3b": "stablelm",
+    "stabilityai/stablelm-tuned-alpha-7b": "stablelm",
 }
 
 DTYPE_DICT = {

generator/__init__.py

@@ -1,3 +1,4 @@
+import os
 import importlib
 import inspect
 
@@ -8,7 +9,11 @@ def choose_bot(config):
     classes = inspect.getmembers(mod, inspect.isclass)
     name, bot_cls = None, None
     for name, bot_cls in classes:
-        if issubclass(bot_cls, ChatBOT):
+        _, filename = os.path.split(inspect.getsourcefile(bot_cls))
+        file_mod, _ = os.path.splitext(filename)
+        # bot_cls may be class that is imported from other files
+        # ex. ChatBOT
+        if file_mod == config.type and issubclass(bot_cls, ChatBOT):
             break
 
     print(f"Choose ChatBOT: {name}")

generator/baize.py

@@ -1,13 +1,11 @@
 import torch
-from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
-from accelerate import init_empty_weights
+from transformers import LlamaForCausalLM
 try:
     from peft import PeftModel
 except:
     PeftModel = None
 
 from .transformersbot import TransformersChatBOT
-from .utils import load_checkpoint_and_dispatch_from_s3
 
 class BaizeBOT(TransformersChatBOT):
     def __init__(self, config):
@@ -17,7 +15,7 @@ def __init__(self, config):
             )
         if config.base_model is None:
             raise ValueError(
-                "Base model's path of Baize should be set."
+                "Base model(llama)'s path of Baize should be set."
             )
         super(BaizeBOT, self).__init__(config)
 
@@ -115,14 +113,8 @@ def process_response(self, response):
             response = response[: response.index("[|Human|]")].strip()
         if "[|AI|]" in response:
             response = response[: response.index("[|AI|]")].strip()
-        
-        return response.strip()
-    
-    def load_tokenizer(self):
-        self.tokenizer = LlamaTokenizer.from_pretrained(
-            self.config.tokenizer_path
-        )
-    
+        return response.strip(" ")
+
     def load_model(self):
 
         llama = self.model_cls.from_pretrained(
@@ -139,6 +131,9 @@ def load_from_s3(self):
         import io
         import json
         from petrel_client.client import Client
+        from accelerate import init_empty_weights
+        from transformers import LlamaConfig
+        from .utils import load_checkpoint_and_dispatch_from_s3
         client = Client()
 
         # get config

generator/belle.py

@@ -1,9 +1,7 @@
 import torch
-from transformers import BloomForCausalLM, AutoConfig
-from accelerate import init_empty_weights
+from transformers import BloomForCausalLM
 
 from .transformersbot import TransformersChatBOT
-from .utils import load_checkpoint_and_dispatch_from_s3
 
 class BELLEBOT(TransformersChatBOT):
     def __init__(self, config):
@@ -42,7 +40,9 @@ def load_from_s3(self):
         import io
         import json
         from petrel_client.client import Client
-        from tqdm import tqdm
+        from accelerate import init_empty_weights
+        from transformers import AutoConfig
+        from .utils import load_checkpoint_and_dispatch_from_s3
         client = Client()
 
         # get model_index

generator/stablelm.py

@@ -0,0 +1,45 @@
+from transformers import GPTNeoXForCausalLM, StoppingCriteria, StoppingCriteriaList
+
+from .transformersbot import TransformersChatBOT
+
+class StableLMBOT(TransformersChatBOT):
+    def __init__(self, config):
+        super(StableLMBOT, self).__init__(config)
+
+    @property
+    def model_cls(self):
+        return GPTNeoXForCausalLM
+    
+    def extra_settings(self):
+        return {
+            "stopping_criteria": StoppingCriteriaList([StopOnTokens()])
+        }
+    
+    def get_prompt(self, query):
+        prompt = """<|SYSTEM|># StableLM Tuned (Alpha version)
+- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
+- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
+- StableLM will refuse to participate in anything that could harm a human.
+"""
+        prompt_dict = {
+            "BOT": "<|ASSISTANT|>{}",
+            "HUMAN": "<|USER|>{}"
+        }
+        for q in query:
+            prompt += prompt_dict[q["role"]].format(q["content"])
+        prompt += "<|ASSISTANT|>"
+
+        return prompt
+    
+    @property
+    def no_split_module_classes(self):
+        return ["GPTNeoXLayer"]
+    
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids, scores, **kwargs) -> bool:
+        stop_ids = [50278, 50279, 50277, 1, 0]
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False

generator/transformersbot.py

@@ -3,10 +3,8 @@
 import torch
 from transformers import AutoTokenizer, AutoConfig
 from transformers.models.auto.modeling_auto import _BaseAutoModelClass
-from accelerate import init_empty_weights
 
 from .chatbot import ChatBOT
-from .utils import load_checkpoint_and_dispatch_from_s3
 
 class TransformersChatBOT(ChatBOT):
     """
@@ -115,6 +113,8 @@ def load_from_s3(self):
         import io
         import json
         from petrel_client.client import Client
+        from accelerate import init_empty_weights
+        from .utils import load_checkpoint_and_dispatch_from_s3
         client = Client()
 
         # get model_index

requirements.txt

@@ -0,0 +1,4 @@
+fastapi
+uvicorn
+transformers
+accelerate