Update app.py

app.py

@@ -1,330 +1,329 @@
 import os
 import torch
 import psutil
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
-from peft import PeftModel, PeftConfig
 from pathlib import Path
-from tqdm import tqdm
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
 from huggingface_hub import login, create_repo, HfApi
-import subprocess
-import math
 import gradio as gr
-import threading
 import queue
 import time
 
-# 创建一个队列用于存储日志消息
+# 全局日志
 log_queue = queue.Queue()
 current_logs = []
 
-
 def log(msg):
-    """统一的日志处理函数"""
+    """追加并打印日志信息"""
     print(msg)
     current_logs.append(msg)
     return "\n".join(current_logs)
 
+def timeit(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        log(f"{func.__name__}: {end_time - start_time:.2f} s")
+        return result   
+    return wrapper
+
+@timeit
 def get_model_size_in_gb(model_name):
-    """估算模型大小（以GB为单位）"""
+    """通过 Hugging Face Hub 元数据估算模型大小（GB）"""
     try:
-        # get model size from huggingface
         api = HfApi()
         model_info = api.model_info(model_name)
+        # 使用 safetensors 大小（不假定文件扩展名）
         return model_info.safetensors.total / (1024 ** 3)
-        
     except Exception as e:
-        log(f"无法估算模型大小: {str(e)}")
-        return 1  # bypass memory check
+        log(f"Unable to estimate model size: {e}")
+        return 1  # 默认值
 
+@timeit
 def check_system_resources(model_name):
-    """检查系统资源并决定使用什么设备"""
-    log("正在检查系统资源...")
-    
-    # 获取系统内存信息
+    """检查系统资源，决定使用 CPU 或 GPU"""
+    log("Checking system resources...")
     system_memory = psutil.virtual_memory()
     total_memory_gb = system_memory.total / (1024 ** 3)
     available_memory_gb = system_memory.available / (1024 ** 3)
+    log(f"Total system memory: {total_memory_gb:.1f}GB")
+    log(f"Available memory: {available_memory_gb:.1f}GB")
     
-    log(f"系统总内存: {total_memory_gb:.1f}GB")
-    log(f"可用内存: {available_memory_gb:.1f}GB")
-    
-    # 估算模型所需内存
     model_size_gb = get_model_size_in_gb(model_name)
-    required_memory_gb = model_size_gb * 2.5  # 需要额外的内存用于计算
-    log(f"估计模型需要内存: {required_memory_gb:.1f}GB")
+    required_memory_gb = model_size_gb * 2.5  # 预留额外内存
+    log(f"Estimated required memory for model: {required_memory_gb:.1f}GB")
     
-    # 检查CUDA是否可用
     if torch.cuda.is_available():
         gpu_name = torch.cuda.get_device_name(0)
         gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
-        log(f"发现GPU: {gpu_name}")
-        log(f"GPU显存: {gpu_memory_gb:.1f}GB")
-        
+        log(f"Detected GPU: {gpu_name} with {gpu_memory_gb:.1f}GB memory")
         if gpu_memory_gb >= required_memory_gb:
-            log("✅ GPU显存足够，将使用GPU进行转换")
+            log("✅ Sufficient GPU memory available; using GPU.")
             return "cuda", gpu_memory_gb
         else:
-            log(f"⚠️ GPU显存不足 (需要 {required_memory_gb:.1f}GB, 实际 {gpu_memory_gb:.1f}GB)")
+            log(f"⚠️ Insufficient GPU memory (requires {required_memory_gb:.1f}GB, found {gpu_memory_gb:.1f}GB).")
     else:
-        log("❌ 未检测到可用的GPU")
-    
-    # 检查CPU内存是否足够
+        log("❌ No GPU detected.")
+        
     if available_memory_gb >= required_memory_gb:
-        log("✅ CPU内存足够，将使用CPU进行转换")
+        log("✅ Sufficient CPU memory available; using CPU.")
         return "cpu", available_memory_gb
     else:
-        raise MemoryError(f"❌ 系统内存不足 (需要 {required_memory_gb:.1f}GB, 可用 {available_memory_gb:.1f}GB)")
+        raise MemoryError(f"❌ Insufficient system memory (requires {required_memory_gb:.1f}GB, available {available_memory_gb:.1f}GB).")
 
+@timeit
 def setup_environment(model_name):
-    # # 检查系统资源并决定使用什么设备
-    # device, available_memory = check_system_resources(model_name)
-    device = "cpu"
+    """选择模型转换时使用的设备"""
+    try:
+        device, _ = check_system_resources(model_name)
+    except Exception as e:
+        log(f"Resource check failed: {e}. Defaulting to CPU.")
+        device = "cpu"
     return device
 
+@timeit
 def create_hf_repo(repo_name, private=True):
-    """创建HuggingFace仓库"""
+    """创建 Hugging Face 仓库（如果不存在的话）"""
     try:
-        # check if repo already exists
         api = HfApi()
+        # 如果仓库已存在，则尝试附加索引直到名称可用
         if api.repo_exists(repo_name):
-            log(f"仓库已存在: {repo_name}")
-            return ValueError(f"仓库已���在: {repo_name}, 请使用其他名称或删除已存在的仓库")
+            retry_index = 0
+            repo_name_with_index = repo_name
+            while api.repo_exists(repo_name_with_index):
+                retry_index += 1
+                log(f"Repository {repo_name_with_index} exists; trying {repo_name}_{retry_index}")
+                repo_name_with_index = f"{repo_name}_{retry_index}"
+            repo_name = repo_name_with_index
         repo_url = create_repo(repo_name, private=private)
-        log(f"创建仓库成功: {repo_url}")
-        return repo_url
+        log(f"Repository created successfully: {repo_url}")
+        return repo_name
     except Exception as e:
-        log(f"创建仓库失败: {str(e)}")
+        log(f"Failed to create repository: {e}")
         raise
 
+@timeit
 def download_and_merge_model(base_model_name, lora_model_name, output_dir, device):
-    log(f"正在加载基础模型: {base_model_name}")
+    """
+    1. 先加载 adapter 的 tokenizer 获取其词表大小
+    2. 加载 base tokenizer 用于后续合并词表
+    3. 加载 base 模型，并将嵌入层调整至 adapter 词表大小
+    4. 使用高层 API 加载 LoRA adapter 并合并其权重
+    5. 求 base 与 adapter tokenizer 的词表并取并集，扩展 tokenizer
+    6. 调整合并模型嵌入层尺寸并保存
+    """
+    model = AutoModelForCausalLM.from_pretrained(base_model_name, low_cpu_mem_usage=True)
+    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))
+    model.load_adapter(lora_model_name, low_cpu_mem_usage=True)
+    model = model.merge_and_unload()
+    model.save_pretrained(output_dir)
+    adapter_tokenizer.save_pretrained(output_dir)
+    return output_dir
+
+@timeit
+def clone_llamacpp_and_download_build():
+    """克隆 llama.cpp 并下载最新构建"""
+    llamacpp_repo = "https://github.com/ggerganov/llama.cpp.git"
+    llamacpp_dir = os.path.join(os.getcwd(), "llama.cpp")
+    
+    if not os.path.exists(llamacpp_dir):
+        log(f"Cloning llama.cpp from {llamacpp_repo}...")
+        os.system(f"git clone {llamacpp_repo} {llamacpp_dir}")
+    
+    log("Building llama.cpp...")
+    build_dir = os.path.join(llamacpp_dir, "build")
+    os.makedirs(build_dir, exist_ok=True)
+    
+    """
+    cmake -B build
+    cmake --build build --config Release
+    """
+    
+    # 进入构建目录并执行 cmake 和 make
+    os.chdir(build_dir)
+    os.system("cmake -B build")
+    os.system("cmake --build build --config Release")
+    
+    log("llama.cpp build completed.")
+    # 返回到原始目录
+    os.chdir(os.path.dirname(llamacpp_dir))
     
-    try:
-        # 先加载原始模型
-        base_model = AutoModelForCausalLM.from_pretrained(
-            base_model_name,
-            torch_dtype=torch.float16,
-            device_map={"": device}
-        )
-        
-        old_vocab_size = base_model.get_input_embeddings().weight.shape[0]
-        print(f"原始词表大小: {old_vocab_size}")
-        # 加载tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-        
-        new_vocab_size = tokenizer.vocab_size
-        print(f"调整词表大小: {old_vocab_size} -> {new_vocab_size}")
-        
-        # 保存原始权重
-        old_embeddings = base_model.get_input_embeddings().weight.data.clone()
-        old_lm_head = base_model.lm_head.weight.data.clone()
-        
-        # 调整词表大小
-        base_model.resize_token_embeddings(new_vocab_size)
-        
-        # 复制原始权重到新的张量
-        with torch.no_grad():
-            base_model.get_input_embeddings().weight.data[:new_vocab_size] = old_embeddings[:new_vocab_size]
-            base_model.lm_head.weight.data[:new_vocab_size] = old_lm_head[:new_vocab_size]
-        
-        log(f"正在加载LoRA模型: {lora_model_name}")
-        log("基础模型配置:" + str(base_model.config))
-        
-        # 加载adapter配置
-        adapter_config = PeftConfig.from_pretrained(lora_model_name)
-        log("Adapter配置:" + str(adapter_config))
-        
-        model = PeftModel.from_pretrained(base_model, lora_model_name)
-        log("正在合并LoRA权重")
-        model = model.merge_and_unload()
 
-        # 创建输出目录
-        output_path = Path(output_dir)
-        output_path.mkdir(parents=True, exist_ok=True)
+@timeit
+def quantize_and_push_model(model_path, repo_id, quant_method=None):
+    """
+    利用 llama-cpp-python 对模型进行量化，并上传到 Hugging Face Hub。
+    使用的量化预设：
+      - 8-bit:  Q8_0
+      - 4-bit:  Q4_K_M 或 Q4_K_L
+      - 2-bit:  Q2_K_L
+    模型输入（model_path）应为全精度（例如 fp16）的 GGUF 文件，
+    输出文件将保存为 <model_path>_q{bits}_{quant_method}
+    """
+    # 使用llama.cpp的转换工具
+    llamacpp_dir = os.path.join(os.getcwd(), "llama.cpp")
+    if not os.path.exists(llamacpp_dir):
+        clone_llamacpp_and_download_build()
 
-        # 保存合并后的模型
-        log(f"正在保存合并后的模型到: {output_dir}")
-        model.save_pretrained(output_dir)
-        tokenizer.save_pretrained(output_dir)
-        
-        
-        return output_dir
-        
-    except Exception as e:
-        log(f"错误: {str(e)}")
-        log(f"错误类型: {type(e)}")
-        import traceback
-        log("详细错误信息:")
-        log(traceback.format_exc())
-        raise
+    # 确保 model_output 目录存在
+    model_output_dir = f"{model_path}/quantized/"
+    os.makedirs(model_output_dir, exist_ok=True)
 
-def quantize_and_push_model(model_path, repo_id, bits=8):
-    """量化模型并推送到HuggingFace"""
-    try:
-        from optimum.bettertransformer import BetterTransformer
-        from transformers import AutoModelForCausalLM
-        
-        log(f"正在加载模型用于{bits}位量化...")
-        model = AutoModelForCausalLM.from_pretrained(model_path)
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-        
-        # 转换为BetterTransformer格式
-        model = BetterTransformer.transform(model)
-        
-        # 量化
-        if bits == 8:
-            from transformers import BitsAndBytesConfig
-            quantization_config = BitsAndBytesConfig(
-                load_in_8bit=True,
-                llm_int8_threshold=6.0
-            )
-        elif bits == 4:
-            from transformers import BitsAndBytesConfig
-            quantization_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.float16,
-                bnb_4bit_quant_type="nf4"
-            )
-        else:
-            raise ValueError(f"不支持的量化位数: {bits}")
-        
-        # 保存量化后的模型
-        quantized_model_path = f"{model_path}_q{bits}"
-        model.save_pretrained(
-            quantized_model_path,
-            quantization_config=quantization_config
-        )
-        tokenizer.save_pretrained(quantized_model_path)
-        
-        # 推送到HuggingFace
-        log(f"正在将{bits}位量化模型推送到HuggingFace...")
-        api = HfApi()
-        api.upload_folder(
-            folder_path=quantized_model_path,
-            repo_id=repo_id,
-            repo_type="model"
-        )
-        log(f"{bits}位量化模型上传完成")
-        
-    except Exception as e:
-        log(f"量化或上传过程中出错: {str(e)}")
-        raise
+    # 中间文件保存在 model_output 目录下
+    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格式")
+        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中间文件已存在，跳过转换")
 
-def process_model(base_model, lora_model, repo_name, hf_token, progress=gr.Progress()):
-    """处理模型的主函数，用于Gradio界面"""
+    # 最终文件保存在 model_output 目录下
+    final_path = os.path.join(model_output_dir, f"{repo_id}-{quant_method}.gguf")
+    print(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}量化文件已存在，跳过量化")
+        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
+
+@timeit
+def process_model(base_model_name, lora_model_name, repo_name, quant_methods, hf_token):
+    """
+    主处理函数：
+      1. 登录并（必要时）创建 Hugging Face 仓库；
+      2. 设置设备；
+      3. 下载并合并 base 模型与 LoRA adapter；
+      4. 异步上传合并后的模型；
+      5. 同时启动四个量化任务（8-bit、2-bit、4-bit 两种模式）；
+      6. 最后统一等待所有 Future 完成，再返回日志。
+    """
     try:
-        login(hf_token) # 我不理解为什么登录一次不行，非得放到环境变量里
+        current_logs.clear()
+        login(hf_token)
         os.environ["HF_TOKEN"] = hf_token
         api = HfApi(token=hf_token)
         username = api.whoami()["name"]
-        if repo_name == "Auto":
-            repo_name = username + "/" + base_model.split("/")[-1] + "_" + lora_model.split("/")[-1]
-        
-        # 清空之前的日志
-        current_logs.clear()
+        if repo_name.strip().lower() == "auto":
+            repo_name = f"{username}/{base_model_name.split('/')[-1]}_{lora_model_name.split('/')[-1]}"
         
-        # 设置环境和检查资源
-        device = setup_environment(base_model)
+        device = setup_environment(base_model_name)
+        repo_name = create_hf_repo(repo_name)
         
-        # 创建HuggingFace仓库
-        repo_url = create_hf_repo(repo_name)
-        
-        # 设置输出目录
         output_dir = os.path.join(".", "output", repo_name)
+        log("Starting model merge process...")
+        model_path = download_and_merge_model(base_model_name, lora_model_name, output_dir, device)
         
-        progress(0.1, desc="开始模型转换流程...")
-        # 下载并合并模型
-        model_path = download_and_merge_model(base_model, lora_model, output_dir, device)
-        
-        # 推送到HuggingFace
-        log(f"正在将模型推送到HuggingFace...")
-        
-        api.upload_folder(
+        # 异步上传合并后的模型
+        log("Scheduling merged model upload...")
+        future_merge = api.upload_large_folder(
             folder_path=model_path,
             repo_id=repo_name,
-            repo_type="model"
+            repo_type="model",
+            num_workers=4,
+            run_as_future=True
         )
         
-        progress(0.4, desc="开始8位量化...")
-        # 量化并上传模型
-        quantize_and_push_model(model_path, repo_name, bits=8)
+        # 启动量化任务，分别使用四种模式：
+        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...")
         
-        progress(0.7, desc="开始4位量化...")
-        quantize_and_push_model(model_path, repo_name, bits=4)
+        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"全部完成！模型已上传至: https://huggingface.co/{repo_name}"
+        final_message = f"All done! Model uploaded to: https://huggingface.co/{repo_name}"
         log(final_message)
-        progress(1.0, desc="处理完成")
-        
-        # remove hf_token from env
-        os.environ.pop("HF_TOKEN")
-        log("HF_TOKEN已从环境变量中删除")
-        
-        # remove model_path
-        os.remove(model_path)
-        log(f"模型路径已删除: {model_path}")
-        
+        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"处理过程中出错: {str(e)}"
+        error_message = f"Error during processing: {e}"
         log(error_message)
         return "\n".join(current_logs)
 
+@timeit
 def create_ui():
-    """创建Gradio界面"""
-    with gr.Blocks(title="模型转换工具") as app:
+    """创建 Gradio 界面，仅展示日志"""
+    with gr.Blocks(title="Model Merge & Quantization Tool") as app:
         gr.Markdown("""
-        # 🤗 模型转换与量化工具
+        # 🤗 Model Merge and Quantization Tool
         
-        这个工具可以帮助你：
-        1. 合并基础模型和LoRA适配器
-        2. 创建4位和8位量化版本
-        3. 自动上传到HuggingFace Hub
+        This tool merges a base model with a LoRA adapter, creates 8-bit, 4-bit and 2-bit quantized versions
+        (using guff's quantization: Q8_0, Q2_K_L, Q4_K_M, Q4_K_L), and uploads them to the Hugging Face Hub.
         """)
-        
         with gr.Row():
             with gr.Column():
                 base_model = gr.Textbox(
-                    label="基础模型路径",
-                    placeholder="例如: Qwen/Qwen2.5-14B-Instruct",
+                    label="Base Model Path",
+                    placeholder="e.g., Qwen/Qwen2.5-14B-Instruct",
                     value="Qwen/Qwen2.5-7B-Instruct"
                 )
                 lora_model = gr.Textbox(
-                    label="LoRA模型路径",
-                    placeholder="输入你的LoRA模型路径"
+                    label="LoRA Model Path",
+                    placeholder="Enter the path to your LoRA model"
                 )
                 repo_name = gr.Textbox(
-                    label="HuggingFace仓库名称",
-                    placeholder="输入要创建的仓库名称",
+                    label="Hugging Face Repository Name",
+                    placeholder="Enter the repository name to create",
                     value="Auto"
                 )
+                quant_method = gr.CheckboxGroup(
+                    choices=["Q2_K", "Q4_K", "IQ4_NL", "Q5_K_M", "Q6_K", "Q8_0"],
+                    value=["Q4_K", "Q8_0"],
+                    label="Quantization Method"
+                )
                 hf_token = gr.Textbox(
-                    label="HuggingFace Token",
-                    placeholder="输入你的HuggingFace Token",
+                    label="Hugging Face Token",
+                    placeholder="Enter your Hugging Face Token",
                     value=os.getenv("HF_TOKEN")
                 )
-                convert_btn = gr.Button("开始转换", variant="primary")
-            
+                convert_btn = gr.Button("Start Conversion", variant="primary")
             with gr.Column():
                 output = gr.TextArea(
-                    label="处理日志",
-                    placeholder="处理日志将在这里显示...",
+                    label="Logs",
+                    placeholder="Processing logs will appear here...",
                     interactive=False,
                     autoscroll=True,
                     lines=20
                 )
-        
-        # 设置事件处理
         convert_btn.click(
             fn=process_model,
-            inputs=[base_model, lora_model, repo_name, hf_token],
+            inputs=[base_model, lora_model, repo_name, quant_method, hf_token],
             outputs=output
         )
-    
     return app
 
+
 if __name__ == "__main__":
-    # 创建并启动Gradio界面
     app = create_ui()
     app.queue()
-    app.launch()
+    app.launch()