Update app.py

app.py

@@ -26,11 +26,11 @@ OUTPUT_DIR = "optimized_models"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 # --- LLAMA.CPP SETUP ---
-## FIX: Define paths at the global scope so all functions can access them.
 LLAMA_CPP_DIR = Path("llama.cpp")
 LLAMA_CPP_CONVERT_SCRIPT = LLAMA_CPP_DIR / "convert.py"
 LLAMA_CPP_QUANTIZE_SCRIPT = LLAMA_CPP_DIR / "quantize" # This is a compiled binary
 
+## FIXED FUNCTION: Replaced 'make' with 'cmake' for the build process.
 def setup_llama_cpp():
     """Clones and builds llama.cpp if not already present."""
     if not LLAMA_CPP_DIR.exists():
@@ -43,14 +43,25 @@ def setup_llama_cpp():
             logging.error(error_msg, exc_info=True)
             raise RuntimeError(error_msg)
 
+    # If the binary doesn't exist, try to build it with CMake.
     if not LLAMA_CPP_QUANTIZE_SCRIPT.exists():
-        logging.info("llama.cpp 'quantize' binary not found. Attempting to build...")
+        logging.info("llama.cpp 'quantize' binary not found. Building with CMake...")
         try:
-            # Use 'make' to build the necessary tools
-            subprocess.run(["make", "-C", str(LLAMA_CPP_DIR), "quantize"], check=True, capture_output=True, text=True)
-            logging.info("'quantize' binary built successfully.")
+            # Step 1: Configure the build directory
+            subprocess.run(
+                ["cmake", "."],
+                cwd=str(LLAMA_CPP_DIR),  # Run command inside the llama.cpp directory
+                check=True, capture_output=True, text=True
+            )
+            # Step 2: Build the 'quantize' target
+            subprocess.run(
+                ["cmake", "--build", ".", "--target", "quantize"],
+                cwd=str(LLAMA_CPP_DIR),
+                check=True, capture_output=True, text=True
+            )
+            logging.info("'quantize' binary built successfully with CMake.")
         except subprocess.CalledProcessError as e:
-            error_msg = f"Failed to build llama.cpp 'quantize' binary. Error: {e.stderr}"
+            error_msg = f"Failed to build llama.cpp with CMake. Error: {e.stderr}"
             logging.error(error_msg, exc_info=True)
             raise RuntimeError(error_msg)
 
@@ -79,7 +90,6 @@ def stage_1_analyze_model(model_id: str):
         logging.error(error_msg)
         return log_stream + error_msg, "Could not analyze model.", gr.Accordion(open=False)
 
-## RE-INTEGRATED: This function is brought back from your original code.
 def stage_2_prune_model(model, prune_percentage: float):
     if prune_percentage == 0:
         return model, "Skipped pruning as percentage was 0."
@@ -108,7 +118,6 @@ def stage_3_4_onnx_quantize(model_path_or_id: str, onnx_quant_type: str, calibra
         raise RuntimeError(f"Failed during `optimum-cli export onnx`. Error:\n{e.stderr}")
 
     try:
-        # For simplicity and stability on HF Spaces, we will only use Dynamic Quantization via CLI.
         quantizer = ORTQuantizer.from_pretrained(onnx_path)
         log_stream += "Performing DYNAMIC quantization...\n"
         dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)
@@ -116,8 +125,6 @@ def stage_3_4_onnx_quantize(model_path_or_id: str, onnx_quant_type: str, calibra
         quantizer.quantize(save_dir=quantized_path, quantization_config=dqconfig)
         log_stream += f"Successfully quantized model to: {quantized_path}\n"
         
-        # If the original input was a model_id, we need to save a new tokenizer.
-        # If it was a local path (from pruning), the tokenizer is already there.
         if not os.path.exists(os.path.join(quantized_path, 'tokenizer_config.json')):
              AutoTokenizer.from_pretrained(model_path_or_id, trust_remote_code=True).save_pretrained(quantized_path)
              log_stream += "Saved new tokenizer files.\n"
@@ -129,7 +136,7 @@ def stage_3_4_onnx_quantize(model_path_or_id: str, onnx_quant_type: str, calibra
 def stage_3_4_gguf_quantize(model_path_or_id: str, original_model_id: str, quantization_strategy: str):
     log_stream = "[STAGE 3 & 4] Converting to GGUF using llama.cpp...\n"
     run_id = datetime.now().strftime("%Y%m%d-%H%M%S")
-    model_name = original_model_id.replace('/', '_') # Use original ID for consistent naming
+    model_name = original_model_id.replace('/', '_')
     gguf_path = os.path.join(OUTPUT_DIR, f"{model_name}-{run_id}-gguf")
     os.makedirs(gguf_path, exist_ok=True)
     
@@ -163,7 +170,6 @@ def stage_3_4_gguf_quantize(model_path_or_id: str, original_model_id: str, quant
         raise RuntimeError(f"An unexpected error occurred during GGUF conversion. Error: {e}")
 
 def stage_5_package_and_upload(model_id: str, optimized_model_path: str, pipeline_log: str, options: dict):
-    # This function is correct from your original version
     log_stream = "[STAGE 5] Packaging and Uploading...\n"
     if not HF_TOKEN:
         return "Skipping upload: HF_TOKEN not found.", log_stream + "Skipping upload: HF_TOKEN not found."
@@ -180,7 +186,6 @@ def stage_5_package_and_upload(model_id: str, optimized_model_path: str, pipelin
     except Exception as e:
         raise RuntimeError(f"Failed to upload to the Hub. Error: {e}")
 
-## RE-INTEGRATED: The main pipeline function now handles both pruning and no-pruning paths.
 def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_percent: float, onnx_quant_type: str, calibration_file, gguf_quant_type: str):
     if not model_id:
         yield {log_output: "Please enter a Model ID.", final_output: "Idle"}
@@ -191,14 +196,13 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
     
     full_log = initial_log
     temp_model_dir = None
-    model_path_or_id = model_id  # Default to memory-efficient path
+    model_path_or_id = model_id
 
     try:
         whoami = api.whoami(token=HF_TOKEN)
         if not whoami: raise RuntimeError("Could not authenticate with Hugging Face Hub. Check your HF_TOKEN.")
         repo_id_for_link = f"{whoami['name']}/{model_id.split('/')[-1]}-amop-cpu-{pipeline_type.lower()}"
         
-        # --- STAGE 2: OPTIONAL PRUNING (Memory-intensive) ---
         if do_prune and prune_percent > 0:
             full_log += f"\n[WARNING] Pruning is memory-intensive and may fail for large models.\n"
             full_log += "Loading base model for pruning...\n"; yield {final_output: "Loading model (1/5)", log_output: full_log}
@@ -210,16 +214,14 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
             model, log = stage_2_prune_model(model, prune_percent)
             full_log += log
             
-            # Save pruned model to a temporary directory for the next stage
             temp_model_dir = tempfile.mkdtemp()
             model.save_pretrained(temp_model_dir)
             tokenizer.save_pretrained(temp_model_dir)
-            model_path_or_id = temp_model_dir # Next stages will use this local path
+            model_path_or_id = temp_model_dir
             full_log += f"Saved intermediate pruned model to {temp_model_dir}\n"
         else:
             full_log += "Pruning skipped.\n"
 
-        # --- STAGE 3 & 4: CONVERSION & QUANTIZATION ---
         if pipeline_type == "ONNX":
             full_log += "Converting to ONNX...\n"; yield {final_output: "Converting to ONNX (3/5)", log_output: full_log}
             optimized_path, log = stage_3_4_onnx_quantize(model_path_or_id, onnx_quant_type, calibration_file.name if onnx_quant_type == "Static" and calibration_file else None)
@@ -232,7 +234,6 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
             raise ValueError("Invalid pipeline type selected.")
         full_log += log
         
-        # --- STAGE 5: UPLOAD ---
         full_log += "Packaging & Uploading...\n"; yield {final_output: "Packaging & Uploading (4/5)", log_output: full_log}
         final_message, log = stage_5_package_and_upload(model_id, optimized_path, full_log, options)
         full_log += log
@@ -243,7 +244,6 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
         full_log += f"\n[ERROR] Pipeline failed: {e}"
         yield {final_output: gr.update(value="ERROR", label="Status"), log_output: full_log, success_box: gr.Markdown(f"❌ **An error occurred.** Check logs for details.", visible=True), run_button: gr.Button(interactive=True, value="Run Optimization Pipeline", variant="primary"), analyze_button: gr.Button(interactive=True, value="Analyze Model")}
     finally:
-        # Clean up the temporary directory if it was created
         if temp_model_dir and os.path.exists(temp_model_dir):
             shutil.rmtree(temp_model_dir)
 
@@ -259,14 +259,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             with gr.Accordion("⚙️ 2. Configure Optimization", open=False) as optimization_accordion:
                 analysis_report_output = gr.Markdown()
                 pipeline_type_radio = gr.Radio(["ONNX", "GGUF"], label="Select Optimization Pipeline")
-                ## RE-INTEGRATED: Pruning UI elements are back.
                 gr.Warning("Pruning requires high RAM and may fail for models >2B parameters on free Spaces.")
                 prune_checkbox = gr.Checkbox(label="Enable Pruning (Optional)", value=False, info="Removes redundant weights before quantization.")
                 prune_slider = gr.Slider(minimum=0, maximum=90, value=20, step=5, label="Pruning Percentage (%)", visible=True)
                 with gr.Group(visible=False) as onnx_options:
                     gr.Markdown("#### ONNX Options")
                     onnx_quant_radio = gr.Radio(["Dynamic"], label="Quantization Type", value="Dynamic", info="Static quantization via UI is not supported.")
-                    calibration_file_upload = gr.File(visible=False) # Keep element for function signature, but hide
+                    calibration_file_upload = gr.File(visible=False)
                 with gr.Group(visible=False) as gguf_options:
                     gr.Markdown("#### GGUF Options")
                     gguf_quant_dropdown = gr.Dropdown(["q4_k_m", "q5_k_m", "q8_0", "f16"], label="Quantization Strategy", value="q4_k_m")
@@ -283,11 +282,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     pipeline_type_radio.change(fn=update_ui_for_pipeline, inputs=pipeline_type_radio, outputs=[onnx_options, gguf_options])
     analyze_button.click(fn=stage_1_analyze_model, inputs=[model_id_input], outputs=[log_output, analysis_report_output, optimization_accordion])
     
-    ## RE-INTEGRATED: Pruning inputs are now passed to the pipeline function.
     run_button.click(fn=run_amop_pipeline, 
                      inputs=[model_id_input, pipeline_type_radio, prune_checkbox, prune_slider, onnx_quant_radio, calibration_file_upload, gguf_quant_dropdown], 
                      outputs=[run_button, analyze_button, final_output, log_output, success_box])
 
 if __name__ == "__main__":
-    # Use .queue() to handle long-running tasks and prevent timeouts
     demo.queue().launch(debug=True)