Update app.py

app.py

@@ -15,7 +15,7 @@ current_logs = []
 
 def log(msg):
     """追加并打印日志信息"""
-    print(msg)
+    log(msg)
     current_logs.append(msg)
     return "\n".join(current_logs)
 
@@ -113,12 +113,19 @@ def download_and_merge_model(base_model_name, lora_model_name, output_dir, devic
     5. 求 base 与 adapter tokenizer 的词表并取并集，扩展 tokenizer
     6. 调整合并模型嵌入层尺寸并保存
     """
+    log("Loading base model...")
     model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True)
+    log("Loading adapter tokenizer...")
     adapter_tokenizer = AutoTokenizer.from_pretrained(lora_model_name)
-    added_tokens_decoder = adapter_tokenizer.added_tokens_decoder
-    model.resize_token_embeddings(adapter_tokenizer.vocab_size + len(added_tokens_decoder))
+    if adapter_tokenizer.pad_token != model.config.pad_token:
+        log("Resizing token embeddings...")
+        added_tokens_decoder = adapter_tokenizer.added_tokens_decoder
+        model.resize_token_embeddings(adapter_tokenizer.vocab_size + len(added_tokens_decoder))
+    log("Loading LoRA adapter...")
     peft_model = PeftModel.from_pretrained(model, lora_model_name, low_cpu_mem_usage=True)
+    log("Merging and unloading model...")
     model = peft_model.merge_and_unload()
+    log("Saving model...")
     model.save_pretrained(output_dir)
     adapter_tokenizer.save_pretrained(output_dir)
     return output_dir
@@ -153,7 +160,7 @@ def clone_llamacpp_and_download_build():
     
 
 @timeit
-def quantize_and_push_model(model_path, repo_id, quant_method=None):
+def quantize(model_path, repo_id, quant_method=None):
     """
     利用 llama-cpp-python 对模型进行量化，并上传到 Hugging Face Hub。
     使用的量化预设：
@@ -176,37 +183,26 @@ def quantize_and_push_model(model_path, repo_id, quant_method=None):
     temp_gguf_path = os.path.join(model_output_dir, f"{repo_id}-f16.gguf")
     
     if not os.path.exists(temp_gguf_path):
-        print(f"正在将模型转换为GGML格式")
+        log(f"正在将模型转换为GGML格式")
         convert_script = os.path.join(llamacpp_dir, "convert_hf_to_gguf.py")
         convert_cmd = f"python {convert_script} {model_path} --outfile {temp_gguf_path}"
         os.system(convert_cmd)
     else:
-        print(f"GGML中间文件已存在，跳过转换")
+        log(f"GGML中间文件已存在，跳过转换")
 
     # 最终文件保存在 model_output 目录下
     final_path = os.path.join(model_output_dir, f"{repo_id}-{quant_method}.gguf")
-    print(f"正在进行{quant_method}量化")
+    log(f"正在进行{quant_method}量化")
     quantize_bin = os.path.join(llamacpp_dir, "build", "bin", "llama-quantize") 
     quant_cmd = f"{quantize_bin} {temp_gguf_path} {final_path} {quant_method}"
     
     if not os.path.exists(final_path):
         os.system(quant_cmd)
     else:
-        print(f"{quant_method}量化文件已存在，跳过量化")
+        log(f"{quant_method}量化文件已存在，跳过量化")
         return None
-        
-    # 异步上传量化模型到 Hugging Face Hub
-    api = HfApi()
-    future = api.upload_file(
-        file_path=final_path,
-        repo_id=repo_id,
-        repo_type="model",
-        commit_message=f"Quantized {quant_method}",
-        commit_description=f"Quantized {model_path} with {quant_method}, using llama.cpp -> {quant_cmd} ",
-        run_as_future=True
-    )
-    log(f"量化模型（{quant_method}）上传已安排；已获得 future 对象。")
-    return future
+    
+    return final_path
 
 @timeit
 def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf_token):
@@ -225,6 +221,10 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
         os.environ["HF_TOKEN"] = hf_token
         api = HfApi(token=hf_token)
         username = api.whoami()["name"]
+        
+        if base_model_name.strip().lower() == "auto":
+            adapter_config = PeftConfig.from_pretrained(lora_model_name)
+            base_model_name = adapter_config.base_model_name_or_path
         if repo_name.strip().lower() == "auto":
             repo_name = f"{username}/{base_model_name.split('/')[-1]}_{lora_model_name.split('/')[-1]}"
         
@@ -235,35 +235,18 @@ def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf
         log("Starting model merge process...")
         model_path = download_and_merge_model(base_model_name, lora_model_name, output_dir, device)
         
-        # 异步上传合并后的模型
-        log("Scheduling merged model upload...")
-        future_merge = api.upload_large_folder(
+        
+        # 量化模型
+        for quant_method in quant_methods:
+            quantize(f"{output_dir}/model.gguf", repo_name, bits=8, quant_method=quant_method)
+        
+        # 上传合并后的模型和量化模型
+        api.upload_large_folder(
             folder_path=model_path,
             repo_id=repo_name,
             repo_type="model",
             num_workers=4,
-            run_as_future=True
         )
-        
-        # 启动量化任务，分别使用四种模式：
-        futures = []
-        for quant_method in quant_methods:
-            future = quantize_and_push_model(f"{output_dir}/model.gguf", repo_name, bits=8, quant_method=quant_method)
-            futures.append(future)
-        log("Background uploads are in progress; performing other tasks if needed...")
-        
-        log("Waiting for merged model upload to complete...")
-        future_merge.result()
-        log("Merged model upload completed.")
-
-        for future in futures:
-            future.result()
-            log(f"{future.result().__name__} completed.")
-        
-        final_message = f"All done! Model uploaded to: https://huggingface.co/{repo_name}"
-        log(final_message)
-        os.environ.pop("HF_TOKEN", None)
-        log("HF_TOKEN removed from environment variables.")
         return "\n".join(current_logs)
     except Exception as e:
         error_message = f"Error during processing: {e}"
@@ -285,7 +268,7 @@ def create_ui():
                 base_model = gr.Textbox(
                     label="Base Model Path",
                     placeholder="e.g., Qwen/Qwen2.5-14B-Instruct",
-                    value="Qwen/Qwen2.5-7B-Instruct"
+                    value="Auto"
                 )
                 lora_model = gr.Textbox(
                     label="LoRA Model Path",