Update README.md

README.md

@@ -6,22 +6,55 @@ license: llama2
 ## 多轮对话测试demo
 ```python
 # -- coding: utf-8 --
-# @time :
+# @time : 2024/12/1 16:26
 # @author : shajiu
 # @email : 18810979033@163.com
 # @file : .py
 # @software: pycharm
+
+
 from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
 import torch
+from peft import PeftModel
+
+class ModelUtils(object):
+
+    @classmethod
+    def load_model(cls, model_name_or_path, load_in_4bit=False, adapter_name_or_path=None):
+        # 是否使用4bit量化进行推理
+        if load_in_4bit:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+                llm_int8_threshold=6.0,
+                llm_int8_has_fp16_weight=False,
+            )
+        else:
+            quantization_config = None
+
+        # 加载base model
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path,
+            load_in_4bit=load_in_4bit,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.float16,
+            device_map='auto',
+            quantization_config=quantization_config
+        )
+
+        # 加载adapter
+        if adapter_name_or_path is not None:
+            model = PeftModel.from_pretrained(model, adapter_name_or_path)
 
-import sys
-sys.path.append("../../")
-from component.utils import ModelUtils
+        return model
 
 
-def main():
+def main(model_name_or_path):
     # 使用合并后的模型进行推理
-    model_name_or_path = 'shajiu/Tibetan_Llama2_7B_Mental_Health'
     adapter_name_or_path = None
 
     # 使用base model和adapter进行推理
@@ -95,5 +128,6 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    model_name_or_path = 'E:\models\shajiuTibetan_Llama2_7B_Mental_Health'
+    main(model_name_or_path)
 ```