1

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel, PeftConfig
 from huggingface_hub import login, create_repo, HfApi
 import gradio as gr
@@ -10,6 +10,8 @@ from gradio_log import Log
 import logging
 
 
+
+
 MEMORY = int(os.getenv("MEMORY", 16)[:-2]) # 64Gi
 CPU_CORES = int(os.getenv("CPU_CORES", 4)) # 4
 SPACE_AUTHOR_NAME = os.getenv("SPACE_AUTHOR_NAME", "Steven10429") # str
@@ -122,15 +124,16 @@ def download_and_merge_model(base_model_name, lora_model_name, output_dir, devic
     6. 调整合并模型嵌入层尺寸并保存
     """
     os.makedirs("temp", exist_ok=True)
+    bnb_config = BitsAndBytesConfig(load_in_8bit=True)
     log.info("Loading base model...")
-    model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True, device_map="auto", force_download=True, trust_remote_code=True, torch_dtype=torch.float16, cache_dir="temp")
+    model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True, device_map="auto", force_download=True, trust_remote_code=True, quantization_config=bnb_config, cache_dir="temp")
     log.info("Loading adapter tokenizer...")
     adapter_tokenizer = AutoTokenizer.from_pretrained(lora_model_name, trust_remote_code=True, device_map="auto", force_download=True)
     log.info("Resizing token embeddings...")
     added_tokens_decoder = adapter_tokenizer.added_tokens_decoder
     model.resize_token_embeddings(adapter_tokenizer.vocab_size + len(added_tokens_decoder))
     log.info("Loading LoRA adapter...")
-    peft_model = PeftModel.from_pretrained(model, lora_model_name, low_cpu_mem_usage=True, device_map="auto", force_download=True, trust_remote_code=True, torch_dtype=torch.float16, cache_dir="temp")
+    peft_model = PeftModel.from_pretrained(model, lora_model_name, low_cpu_mem_usage=True, device_map="auto", force_download=True, trust_remote_code=True, quantization_config=bnb_config, cache_dir="temp")
     log.info("Merging and unloading model...")
     model = peft_model.merge_and_unload()
     log.info("Saving model...")
@@ -192,25 +195,25 @@ def quantize(model_path, repo_id, quant_method=None):
     os.makedirs(model_output_dir, exist_ok=True)
 
     # 中间文件保存在 model_output 目录下
-    guff_16_path =f"./{repo_id}-f16.gguf"
+    guff_8_path =f"./{repo_id}-q8_0.gguf"
     
-    if not os.path.exists(guff_16_path):
+    if not os.path.exists(guff_8_path):
         log.info(f"正在将模型转换为GGML格式")
         convert_script = os.path.join(llamacpp_dir, "convert_hf_to_gguf.py")
-        convert_cmd = f"python {convert_script} {model_path} --outfile {guff_16_path}"
+        convert_cmd = f"python {convert_script} {model_path} --outfile {guff_8_path} --outtype q8_0"
         print(f"syscall:[{convert_cmd}]")
         os.system(convert_cmd)
     else:
         log.info(f"GGML中间文件已存在，跳过转换")
         
-    if quant_method == "fp16":
-        return guff_16_path  # for upload to hub
+    if quant_method.lower() == "q8_0":
+        return guff_8_path  # for upload to hub
 
     # 最终文件保存在 model_output 目录下
     final_path = os.path.join(model_output_dir, f"{repo_id}-{quant_method}.gguf")
     log.info(f"正在进行{quant_method}量化")
     quantize_bin = os.path.join(llamacpp_dir, "build", "bin", "llama-quantize") 
-    quant_cmd = f"{quantize_bin} {guff_16_path} {final_path} {quant_method}"
+    quant_cmd = f"{quantize_bin} {guff_8_path} {final_path} {quant_method}"
     print(f"syscall:[{quant_cmd}]")
     
     if not os.path.exists(final_path):
@@ -230,14 +233,12 @@ tags:
 - text-generation
 - peft{quantization}
 library_name: transformers
-base_model: {base_model_name}
+base_model: {base_model_name}{lora_model_name}
 widget:
 - messages:
     - role: user
       content: What is your favorite condiment?
 license: other
-datasets:
-- {lora_model_name}
 ---
 # Model
 
@@ -310,29 +311,29 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
         )
         log.info("Upload completed.")
         
-        # remove model for space limit
+        if len(quant_methods) > 0:
+            quantize(output_dir, repo_name, "Q8_0")
+            # remove model for space limit
         shutil.rmtree(model_path)
-        
+        log.info("Removed model from local")
+            
         os.makedirs(os.path.join(output_dir, "quantized"), exist_ok=True)
         if len(quant_methods) > 0:
-            quantize(output_dir, repo_name, "fp16") # for
-            # 量化模型
             for quant_method in quant_methods:
                 quantize(output_dir, repo_name, quant_method=quant_method)
             os.system(f"mv ./{repo_name}-f16.gguf ./{output_dir}/quantized/")
             
-        api.upload_folder(
-            folder_path=os.path.join(output_dir, "quantized"),
-            path_in_repo="quantized",
-            repo_id=repo_name,
-            repo_type="model",
-            num_workers=os.cpu_count() if os.cpu_count() > 4 else 4,
-            print_report_every=10,
-        )
+            api.upload_large_folder(
+                folder_path=os.path.join(output_dir, "quantized"),
+                repo_id=repo_name,
+                repo_type="model",
+                num_workers=os.cpu_count() if os.cpu_count() > 4 else 4,
+                print_report_every=10,
+            )
         
-        # rm -rf model_path
-        shutil.rmtree(model_path)
-        log.info("Removed model from local")
+            # rm -rf model_path
+            shutil.rmtree(model_path)
+            log.info("Removed model from local")
         
     except Exception as e:
         error_message = f"Error during processing: {e}"