feat: 加入LoRA功能

modules/models/LLaMA.py

@@ -11,10 +11,6 @@ from ..presets import *
 from ..utils import *
 from .base_model import BaseLLMModel
 
-import json
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
-
 SYS_PREFIX = "<<SYS>>\n"
 SYS_POSTFIX = "\n<</SYS>>\n\n"
 INST_PREFIX = "<s>[INST] "
@@ -22,6 +18,7 @@ INST_POSTFIX = " "
 OUTPUT_PREFIX = "[/INST] "
 OUTPUT_POSTFIX = "</s>"
 
+
 def download(repo_id, filename, retry=10):
     if os.path.exists("./models/downloaded_models.json"):
         with open("./models/downloaded_models.json", "r") as f:
@@ -32,7 +29,12 @@ def download(repo_id, filename, retry=10):
         downloaded_models = {}
     while retry > 0:
         try:
-            model_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir="models", resume_download=True)
+            model_path = hf_hub_download(
+                repo_id=repo_id,
+                filename=filename,
+                cache_dir="models",
+                resume_download=True,
+            )
             downloaded_models[repo_id] = {"path": model_path}
             with open("./models/downloaded_models.json", "w") as f:
                 json.dump(downloaded_models, f)
@@ -46,57 +48,69 @@ def download(repo_id, filename, retry=10):
 
 
 class LLaMA_Client(BaseLLMModel):
-    def __init__(
-        self,
-        model_name,
-        lora_path=None,
-        user_name=""
-    ) -> None:
+    def __init__(self, model_name, lora_path=None, user_name="") -> None:
         super().__init__(model_name=model_name, user=user_name)
 
         self.max_generation_token = 1000
-        self.end_string = "\n\n"
-        # We don't need input data
-        path_to_model = download(MODEL_METADATA[model_name]["repo_id"], MODEL_METADATA[model_name]["filelist"][0])
+        if model_name in MODEL_METADATA:
+            path_to_model = download(
+                MODEL_METADATA[model_name]["repo_id"],
+                MODEL_METADATA[model_name]["filelist"][0],
+            )
+        else:
+            dir_to_model = os.path.join("models", model_name)
+            # look for nay .gguf file in the dir_to_model directory and its subdirectories
+            path_to_model = None
+            for root, dirs, files in os.walk(dir_to_model):
+                for file in files:
+                    if file.endswith(".gguf"):
+                        path_to_model = os.path.join(root, file)
+                        break
+                if path_to_model is not None:
+                    break
         self.system_prompt = ""
 
-        global LLAMA_MODEL
-        if LLAMA_MODEL is None:
-            LLAMA_MODEL = Llama(model_path=path_to_model)
-            # 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:
-            #     model_source = f"decapoda-research/{model_name}"
-                # raise Exception(f"models目录下没有这个模型: {model_name}")
-            # if lora_path is not None:
-            #     lora_path = f"lora/{lora_path}"
+        if lora_path is not None:
+            lora_path = os.path.join("lora", lora_path)
+            self.model = Llama(model_path=path_to_model, lora_path=lora_path)
+        else:
+            self.model = Llama(model_path=path_to_model)
 
     def _get_llama_style_input(self):
         context = []
         for conv in self.history:
             if conv["role"] == "system":
-                context.append(SYS_PREFIX+conv["content"]+SYS_POSTFIX)
+                context.append(SYS_PREFIX + conv["content"] + SYS_POSTFIX)
             elif conv["role"] == "user":
-                context.append(INST_PREFIX+conv["content"]+INST_POSTFIX+OUTPUT_PREFIX)
+                context.append(
+                    INST_PREFIX + conv["content"] + INST_POSTFIX + OUTPUT_PREFIX
+                )
             else:
-                context.append(conv["content"]+OUTPUT_POSTFIX)
+                context.append(conv["content"] + OUTPUT_POSTFIX)
         return "".join(context)
 
     def get_answer_at_once(self):
         context = self._get_llama_style_input()
-        response = LLAMA_MODEL(context, max_tokens=self.max_generation_token, stop=[], echo=False, stream=False)
+        response = self.model(
+            context,
+            max_tokens=self.max_generation_token,
+            stop=[],
+            echo=False,
+            stream=False,
+        )
         return response, len(response)
 
     def get_answer_stream_iter(self):
         context = self._get_llama_style_input()
-        iter = LLAMA_MODEL(context, max_tokens=self.max_generation_token, stop=[], echo=False, stream=True)
+        iter = self.model(
+            context,
+            max_tokens=self.max_generation_token,
+            stop=[],
+            echo=False,
+            stream=True,
+        )
         partial_text = ""
         for i in iter:
             response = i["choices"][0]["text"]
             partial_text += response
-            yield partial_text
+            yield partial_text

modules/models/models.py

@@ -26,7 +26,7 @@ def get_model(
     msg = i18n("模型设置为了：") + f" {model_name}"
     model_type = ModelType.get_type(model_name)
     lora_selector_visibility = False
-    lora_choices = []
+    lora_choices = ["No LoRA"]
     dont_change_lora_selector = False
     if model_type != ModelType.OpenAI:
         config.local_embedding = True
@@ -55,8 +55,7 @@ def get_model(
             logging.info(msg)
             lora_selector_visibility = True
             if os.path.isdir("lora"):
-                get_file_names_by_pinyin("lora", filetypes=[""])
-            lora_choices = ["No LoRA"] + lora_choices
+                lora_choices = ["No LoRA"] + get_file_names_by_pinyin("lora", filetypes=[""])
         elif model_type == ModelType.LLaMA and lora_model_path != "":
             logging.info(f"正在加载LLaMA模型: {model_name} + {lora_model_path}")
             from .LLaMA import LLaMA_Client