Update app.py

app.py

@@ -10,8 +10,8 @@ from huggingface_hub import HfApi
 from transformers import AutoConfig, AutoModel, AutoTokenizer
 from optimum.onnxruntime import ORTQuantizer, ORTModelForCausalLM
 from optimum.onnxruntime.configuration import AutoQuantizationConfig
-from optimum.exporters.onnx import main_export as onnx_export
-from optimum.exporters.gguf import main_export as gguf_export
+# Use the unified optimum.main_export entrypoint
+from optimum.exporters.main import main_export
 import torch.nn.utils.prune as prune
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -67,7 +67,7 @@ def stage_3_4_onnx_quantize(model_path: str, calibration_data_path: str):
         model_name = os.path.basename(model_path)
         onnx_path = os.path.join(OUTPUT_DIR, f"{model_name}-{run_id}-onnx")
         
-        onnx_export(model_path, output=onnx_path, task="auto", trust_remote_code=True)
+        main_export(model_path, output=onnx_path, task="auto", trust_remote_code=True)
         log_stream += f"Successfully exported base model to ONNX at: {onnx_path}\n"
         
         quantizer = ORTQuantizer.from_pretrained(onnx_path)
@@ -105,7 +105,7 @@ def stage_3_4_gguf_quantize(model_id: str, quantization_strategy: str):
         gguf_path = os.path.join(OUTPUT_DIR, f"{model_name}-{run_id}-gguf")
         os.makedirs(gguf_path, exist_ok=True)
         
-        gguf_export(model_id, output=os.path.join(gguf_path, "model.gguf"), quantization_strategy=quantization_strategy, trust_remote_code=True)
+        main_export(model_id, output=os.path.join(gguf_path, "model.gguf"), export_format="gguf", quantization_strategy=quantization_strategy, trust_remote_code=True)
 
         log_stream += f"Successfully exported and quantized model to GGUF at: {gguf_path}\n"
         return gguf_path, log_stream