Update app.py

app.py

@@ -14,7 +14,6 @@ from optimum.onnxruntime import ORTQuantizer
 from optimum.onnxruntime.configuration import AutoQuantizationConfig
 import torch.nn.utils.prune as prune
 
-# --- SETUP ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -25,23 +24,27 @@ api = HfApi()
 OUTPUT_DIR = "optimized_models"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# --- LLAMA.CPP SETUP ---
 LLAMA_CPP_DIR = Path("llama.cpp")
-## FINAL FIX: The correct, stable script is in the 'tools' subdirectory.
-LLAMA_CPP_CONVERT_SCRIPT = LLAMA_CPP_DIR / "tools" / "convert-hf-to-gguf.py"
+LLAMA_CPP_CONVERT_SCRIPT = LLAMA_CPP_DIR / "convert.py"
 LLAMA_CPP_QUANTIZE_SCRIPT = LLAMA_CPP_DIR / "quantize"
 
 def setup_llama_cpp():
-    """Clones and builds llama.cpp if not already present."""
     if not LLAMA_CPP_DIR.exists():
         logging.info("Cloning llama.cpp repository...")
         try:
             subprocess.run(["git", "clone", "https://github.com/ggerganov/llama.cpp.git"], check=True, capture_output=True, text=True)
             logging.info("llama.cpp cloned successfully.")
         except subprocess.CalledProcessError as e:
-            error_msg = f"Failed to clone llama.cpp. Error: {e.stderr}"
-            logging.error(error_msg, exc_info=True)
-            raise RuntimeError(error_msg)
+            raise RuntimeError(f"Failed to clone llama.cpp. Error: {e.stderr}")
+
+    requirements_path = LLAMA_CPP_DIR / "requirements.txt"
+    if requirements_path.exists():
+        logging.info("Installing llama.cpp Python dependencies...")
+        try:
+            subprocess.run(["pip", "install", "-r", str(requirements_path)], check=True, capture_output=True, text=True)
+            logging.info("llama.cpp Python dependencies installed successfully.")
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(f"Failed to install llama.cpp requirements. Error: {e.stderr}")
 
     if not LLAMA_CPP_QUANTIZE_SCRIPT.exists():
         logging.info("llama.cpp 'quantize' binary not found. Building with CMake...")
@@ -50,17 +53,13 @@ def setup_llama_cpp():
             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 with CMake. Error: {e.stderr}"
-            logging.error(error_msg, exc_info=True)
-            raise RuntimeError(error_msg)
+            raise RuntimeError(f"Failed to build llama.cpp with CMake. Error: {e.stderr}")
 
-# Run setup on script start
 try:
     setup_llama_cpp()
 except Exception as e:
     logging.error(f"FATAL ERROR during llama.cpp setup: {e}", exc_info=True)
 
-
 def stage_1_analyze_model(model_id: str):
     log_stream = "[STAGE 1] Analyzing model...\n"
     try:
@@ -95,17 +94,13 @@ def stage_3_4_onnx_quantize(model_path_or_id: str, onnx_quant_type: str, calibra
     run_id = datetime.now().strftime("%Y%m%d-%H%M%S")
     model_name = model_path_or_id.split('/')[-1]
     onnx_path = os.path.join(OUTPUT_DIR, f"{model_name}-{run_id}-onnx")
-
     try:
-        log_stream += f"Executing `optimum-cli export onnx` for '{model_path_or_id}'...\n"
         export_command = ["optimum-cli", "export", "onnx", "--model", model_path_or_id, "--trust-remote-code", onnx_path]
         process = subprocess.run(export_command, check=True, capture_output=True, text=True)
-        log_stream += process.stdout
+        log_stream += f"Executing `optimum-cli export onnx` for '{model_path_or_id}'...\n{process.stdout}\n"
         if process.stderr: log_stream += f"[STDERR]\n{process.stderr}\n"
-        log_stream += f"Successfully exported to ONNX at: {onnx_path}\n"
     except subprocess.CalledProcessError as e:
         raise RuntimeError(f"Failed during `optimum-cli export onnx`. Error:\n{e.stderr}")
-
     try:
         quantizer = ORTQuantizer.from_pretrained(onnx_path)
         log_stream += "Performing DYNAMIC quantization...\n"
@@ -113,11 +108,9 @@ def stage_3_4_onnx_quantize(model_path_or_id: str, onnx_quant_type: str, calibra
         quantized_path = os.path.join(onnx_path, "quantized-dynamic")
         quantizer.quantize(save_dir=quantized_path, quantization_config=dqconfig)
         log_stream += f"Successfully quantized model to: {quantized_path}\n"
-        
         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"
-             
         return quantized_path, log_stream
     except Exception as e:
         raise RuntimeError(f"Failed during ONNX quantization step. Error: {e}")
@@ -128,20 +121,15 @@ def stage_3_4_gguf_quantize(model_path_or_id: str, original_model_id: str, quant
     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)
-    
     f16_gguf_path = os.path.join(gguf_path, "model-f16.gguf")
     quantized_gguf_path = os.path.join(gguf_path, "model.gguf")
-
     try:
-        log_stream += f"Executing llama.cpp conversion script on '{model_path_or_id}'...\n"
         convert_command = ["python3", str(LLAMA_CPP_CONVERT_SCRIPT), model_path_or_id, "--outfile", f16_gguf_path, "--outtype", "f16"]
         process = subprocess.run(convert_command, check=True, capture_output=True, text=True)
-        log_stream += process.stdout
+        log_stream += f"Executing llama.cpp conversion script on '{model_path_or_id}'...\n{process.stdout}\n"
         if process.stderr: log_stream += f"[STDERR]\n{process.stderr}\n"
-        
         quantize_map = {"q4_k_m": "Q4_K_M", "q5_k_m": "Q5_K_M", "q8_0": "Q8_0", "f16": "F16"}
         target_quant_name = quantize_map.get(quantization_strategy.lower(), "Q4_K_M")
-
         if target_quant_name == "F16":
             log_stream += "Target is F16, renaming file...\n"
             os.rename(f16_gguf_path, quantized_gguf_path)
@@ -149,7 +137,7 @@ def stage_3_4_gguf_quantize(model_path_or_id: str, original_model_id: str, quant
             log_stream += f"Quantizing FP16 GGUF to {target_quant_name}...\n"
             quantize_command = [str(LLAMA_CPP_QUANTIZE_SCRIPT), f16_gguf_path, quantized_gguf_path, target_quant_name]
             process = subprocess.run(quantize_command, check=True, capture_output=True, text=True)
-            log_stream += process.stdout
+            log_stream += f"{process.stdout}\n"
             if process.stderr: log_stream += f"[STDERR]\n{process.stderr}\n"
             os.remove(f16_gguf_path)
         return gguf_path, log_stream
@@ -179,30 +167,24 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
     if not model_id:
         yield {log_output: "Please enter a Model ID.", final_output: "Idle"}
         return
-
     initial_log = f"[START] AMOP {pipeline_type} Pipeline Initiated for '{model_id}'.\n"
     yield {run_button: gr.Button(interactive=False, value="🚀 Running..."), analyze_button: gr.Button(interactive=False), final_output: f"RUNNING ({pipeline_type})", log_output: initial_log}
-    
     full_log = initial_log
     temp_model_dir = None
     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()}"
-        
         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}
+            yield {final_output: "Loading model (1/5)", log_output: full_log}
             model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
             tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
             full_log += f"Successfully loaded '{model_id}'.\n"
-
             yield {final_output: "Pruning model (2/5)", log_output: full_log}
             model, log = stage_2_prune_model(model, prune_percent)
             full_log += log
-            
             temp_model_dir = tempfile.mkdtemp()
             model.save_pretrained(temp_model_dir)
             tokenizer.save_pretrained(temp_model_dir)
@@ -210,23 +192,20 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
             full_log += f"Saved intermediate pruned model to {temp_model_dir}\n"
         else:
             full_log += "Pruning skipped.\n"
-
         if pipeline_type == "ONNX":
-            full_log += "Converting to ONNX...\n"; yield {final_output: "Converting to ONNX (3/5)", log_output: full_log}
+            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)
             options = {'pipeline_type': 'ONNX', 'prune': do_prune, 'prune_percent': prune_percent, 'quant_type': onnx_quant_type}
         elif pipeline_type == "GGUF":
-            full_log += "Converting to GGUF...\n"; yield {final_output: "Converting to GGUF (3/5)", log_output: full_log}
+            yield {final_output: "Converting to GGUF (3/5)", log_output: full_log}
             optimized_path, log = stage_3_4_gguf_quantize(model_path_or_id, model_id, gguf_quant_type)
             options = {'pipeline_type': 'GGUF', 'prune': do_prune, 'prune_percent': prune_percent, 'quant_type': gguf_quant_type}
         else:
             raise ValueError("Invalid pipeline type selected.")
         full_log += log
-        
-        full_log += "Packaging & Uploading...\n"; yield {final_output: "Packaging & Uploading (4/5)", log_output: full_log}
+        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
-
         yield {final_output: gr.update(value="SUCCESS", label="Status"), log_output: full_log, success_box: gr.Markdown(f"✅ **Success!** Model available: [{repo_id_for_link}](https://huggingface.co/{repo_id_for_link})", visible=True), run_button: gr.Button(interactive=True, value="Run Optimization Pipeline", variant="primary"), analyze_button: gr.Button(interactive=True, value="Analyze Model")}
     except Exception as e:
         logging.error(f"AMOP Pipeline failed. Error: {e}", exc_info=True)
@@ -236,7 +215,6 @@ def run_amop_pipeline(model_id: str, pipeline_type: str, do_prune: bool, prune_p
         if temp_model_dir and os.path.exists(temp_model_dir):
             shutil.rmtree(temp_model_dir)
 
-# --- GRADIO UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 AMOP: Adaptive Model Optimization Pipeline")
     if not HF_TOKEN: gr.Warning("HF_TOKEN not set! The final 'upload' step will be skipped.")
@@ -264,13 +242,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             final_output = gr.Label(value="Idle", label="Status")
             success_box = gr.Markdown(visible=False)
             log_output = gr.Textbox(label="Live Logs", lines=20, interactive=False)
-
     def update_ui_for_pipeline(pipeline_type):
         return {onnx_options: gr.Group(visible=pipeline_type=="ONNX"), gguf_options: gr.Group(visible=pipeline_type=="GGUF")}
-    
     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])
-    
     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])