加入了llama模型支持

ChuanhuChatbot.py

@@ -80,6 +80,9 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                     model_select_dropdown = gr.Dropdown(
                         label="选择模型", choices=MODELS, multiselect=False, value=MODELS[DEFAULT_MODEL], interactive=True
                     )
+                    lora_select_dropdown = gr.Dropdown(
+                        label="选择LoRA模型", choices=[], multiselect=False, interactive=True, visible=False
+                    )
                     use_streaming_checkbox = gr.Checkbox(
                         label="实时传输回答", value=True, visible=ENABLE_STREAMING_OPTION
                     )
@@ -350,7 +353,8 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     # LLM Models
     keyTxt.change(current_model.value.set_key, keyTxt, [status_display]).then(**get_usage_args)
     keyTxt.submit(**get_usage_args)
-    model_select_dropdown.change(current_model.value.get_model, [model_select_dropdown, keyTxt, temperature_slider, top_p_slider, systemPromptTxt], [status_display], show_progress=True)
+    model_select_dropdown.change(current_model.value.get_model, [model_select_dropdown, lora_select_dropdown, keyTxt, temperature_slider, top_p_slider, systemPromptTxt], [status_display, lora_select_dropdown], show_progress=True)
+    lora_select_dropdown.change(current_model.value.get_model, [model_select_dropdown, lora_select_dropdown, keyTxt, temperature_slider, top_p_slider, systemPromptTxt], [status_display], show_progress=True)
 
     # Template
     systemPromptTxt.change(current_model.value.set_system_prompt, [systemPromptTxt], None)

configs/ds_config_chatbot.json

@@ -0,0 +1,17 @@
+{
+    "fp16": {
+        "enabled": false
+    },
+    "bf16": {
+        "enabled": true
+    },
+    "comms_logger": {
+        "enabled": false,
+        "verbose": false,
+        "prof_all": false,
+        "debug": false
+    },
+    "steps_per_print": 20000000000000000,
+    "train_micro_batch_size_per_gpu": 1,
+    "wall_clock_breakdown": false
+}

modules/base_model.py

@@ -24,6 +24,7 @@ from .config import retrieve_proxy
 
 
 class ModelType(Enum):
+    Unknown = -1
     OpenAI = 0
     ChatGLM = 1
     LLaMA = 2
@@ -31,12 +32,15 @@ class ModelType(Enum):
     @classmethod
     def get_type(cls, model_name: str):
         model_type = None
-        if "gpt" in model_name.lower():
+        model_name_lower = model_name.lower()
+        if "gpt" in model_name_lower:
             model_type = ModelType.OpenAI
-        elif "chatglm" in model_name.lower():
+        elif "chatglm" in model_name_lower:
             model_type = ModelType.ChatGLM
-        else:
+        elif "llama" in model_name_lower:
             model_type = ModelType.LLaMA
+        else:
+            model_type = ModelType.Unknown
         return model_type
 
 

modules/models.py

@@ -13,11 +13,6 @@ import platform
 from dataclasses import dataclass, field
 from transformers import HfArgumentParser
 
-from lmflow.datasets.dataset import Dataset
-from lmflow.pipeline.auto_pipeline import AutoPipeline
-from lmflow.models.auto_model import AutoModel
-from lmflow.args import ModelArguments, DatasetArguments, AutoArguments
-
 from tqdm import tqdm
 import colorama
 from duckduckgo_search import ddg
@@ -203,12 +198,13 @@ class OpenAIClient(BaseLLMModel):
 
 
 class ChatGLM_Client(BaseLLMModel):
-    def __init__(self, model_name, model_path=None) -> None:
+    def __init__(self, model_name) -> None:
         super().__init__(model_name=model_name)
         from transformers import AutoTokenizer, AutoModel
         import torch
 
         system_name = platform.system()
+        model_path=None
         if os.path.exists("models"):
             model_dirs = os.listdir("models")
             if model_name in model_dirs:
@@ -285,18 +281,49 @@ class LLaMA_Client(BaseLLMModel):
         lora_path=None,
     ) -> None:
         super().__init__(model_name=model_name)
+        from lmflow.datasets.dataset import Dataset
+        from lmflow.pipeline.auto_pipeline import AutoPipeline
+        from lmflow.models.auto_model import AutoModel
+        from lmflow.args import ModelArguments, DatasetArguments, AutoArguments, InferencerArguments
+        model_path = None
+        if os.path.exists("models"):
+            model_dirs = os.listdir("models")
+            if model_name in model_dirs:
+                model_path = f"models/{model_name}"
+        if model_path is not None:
+            model_source = model_path
+        else:
+            raise Exception(f"models目录下没有这个模型: {model_name}")
         self.max_generation_token = 1000
         pipeline_name = "inferencer"
         PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)
 
-        parser = HfArgumentParser(
-            (
-                ModelArguments,
-                PipelineArguments,
-                ChatbotArguments,
-            )
-        )
-        model_args, pipeline_args, chatbot_args = parser.parse_args_into_dataclasses()
+        """
+        if [ $# -ge 2 ]; then
+        lora_args="--lora_model_path $2"
+        fi
+        CUDA_VISIBLE_DEVICES=2 \
+        deepspeed examples/chatbot.py \
+            --deepspeed configs/ds_config_chatbot.json \
+            --model_name_or_path ${model} \
+            ${lora_args}
+
+        model_args:
+        ModelArguments(model_name_or_path='/home/guest/llm_models/llama/7B', lora_model_path='/home/guest/llm_models/lora/baize-lora-7B', model_type=None, config_overrides=None, config_name=None, tokenizer_name=None, cache_dir=None, use_fast_tokenizer=True, model_revision='main', use_auth_token=False, torch_dtype=None, use_lora=False, lora_r=8, lora_alpha=32, lora_dropout=0.1, use_ram_optimized_load=True)
+        pipeline_args:
+        InferencerArguments(local_rank=0, random_seed=1, deepspeed='configs/ds_config_chatbot.json', mixed_precision='bf16')
+        """
+
+        # parser = HfArgumentParser(
+        #     (
+        #         ModelArguments,
+        #         PipelineArguments,
+        #         ChatbotArguments,
+        #     )
+        # )
+        model_args = ModelArguments(model_name_or_path=model_source, lora_model_path=lora_path, model_type=None, config_overrides=None, config_name=None, tokenizer_name=None, cache_dir=None, use_fast_tokenizer=True, model_revision='main', use_auth_token=False, torch_dtype=None, use_lora=False, lora_r=8, lora_alpha=32, lora_dropout=0.1, use_ram_optimized_load=True)
+        pipeline_args = InferencerArguments(local_rank=0, random_seed=1, deepspeed='configs/ds_config_chatbot.json', mixed_precision='bf16')
+        # model_args, pipeline_args, chatbot_args = parser.parse_args_into_dataclasses()
 
         with open(pipeline_args.deepspeed, "r") as f:
             ds_config = json.load(f)
@@ -377,23 +404,54 @@ class ModelManager:
         top_p=None,
         system_prompt=None,
     ) -> BaseLLMModel:
+        print(lora_model_path)
         msg = f"模型设置为了： {model_name}"
         logging.info(msg)
         model_type = ModelType.get_type(model_name)
+        lora_selector_visibility = False
+        lora_choices = []
+        dont_change_lora_selector = False
         if model_type != ModelType.OpenAI:
             config.local_embedding = True
-        if model_type == ModelType.OpenAI:
-            model = OpenAIClient(
-                model_name=model_name,
-                api_key=access_key,
-                system_prompt=system_prompt,
-                temperature=temperature,
-                top_p=top_p,
-            )
-        elif model_type == ModelType.ChatGLM:
-            model = ChatGLM_Client(model_name)
-        self.model = model
-        return msg
+        model = None
+        try:
+            if model_type == ModelType.OpenAI:
+                model = OpenAIClient(
+                    model_name=model_name,
+                    api_key=access_key,
+                    system_prompt=system_prompt,
+                    temperature=temperature,
+                    top_p=top_p,
+                )
+            elif model_type == ModelType.ChatGLM:
+                model = ChatGLM_Client(model_name)
+            elif model_type == ModelType.LLaMA and lora_model_path == "":
+                msg = "现在请选择LoRA模型"
+                logging.info(msg)
+                lora_selector_visibility = True
+                if os.path.isdir("lora"):
+                    lora_choices = get_file_names("lora", plain=True, filetypes=[""])
+                lora_choices = ["No LoRA"] + lora_choices
+            elif model_type == ModelType.LLaMA and lora_model_path != "":
+                dont_change_lora_selector = True
+                if lora_model_path == "No LoRA":
+                    lora_model_path = None
+                    msg += " + No LoRA"
+                else:
+                    msg += f" + {lora_model_path}"
+                model = LLaMA_Client(model_name, lora_model_path)
+                pass
+            elif model_type == ModelType.Unknown:
+                raise ValueError(f"未知模型: {model_name}")
+        except Exception as e:
+            logging.error(e)
+            msg = f"{STANDARD_ERROR_MSG}: {e}"
+        if model is not None:
+            self.model = model
+        if dont_change_lora_selector:
+            return msg
+        else:
+            return msg, gr.Dropdown.update(choices=lora_choices, visible=lora_selector_visibility)
 
     def predict(self, *args):
         iter = self.model.predict(*args)

modules/presets.py

@@ -59,7 +59,15 @@ MODELS = [
     "gpt-4-32k-0314",
     "chatglm-6b",
     "chatglm-6b-int4",
-    "chatglm-6b-int4-qe"
+    "chatglm-6b-int4-qe",
+    "llama-7b-hf",
+    "llama-7b-hf-int4",
+    "llama-7b-hf-int8",
+    "llama-13b-hf",
+    "llama-13b-hf-int4",
+    "llama-30b-hf",
+    "llama-30b-hf-int4",
+    "llama-65b-hf",
 ]  # 可选的模型
 
 DEFAULT_MODEL = 0  # 默认的模型在MODELS中的序号，从0开始数

modules/utils.py

@@ -445,3 +445,7 @@ def get_last_day_of_month(any_day):
     next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
     # subtracting the number of the current day brings us back one month
     return next_month - datetime.timedelta(days=next_month.day)
+
+def get_model_source(model_name, alternative_source):
+    if model_name == "gpt2-medium":
+        return "https://huggingface.co/gpt2-medium"