app.py

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
import gradio as gr
import logging
from typing import Tuple, Dict, Any
from src.utilities.resources import ResourceManager
from src.utilities.push_to_hub import push_to_hub
from src.optimizations.onnx_conversion import convert_to_onnx
from src.optimizations.quantize import quantize_onnx_model
from src.handlers import get_model_handler, TASK_CONFIGS
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
import json

def process_model(

    model_name: str,

    task: str,

    quantization_type: str,

    enable_onnx: bool,

    onnx_quantization: str,

    hf_token: str,

    repo_name: str,

    test_text: str

) -> Tuple[Dict[str, Any], str, Dict[str, Any]]:
    try:
        resource_manager = ResourceManager()
        status_updates = []
        status = {
            "status": "Processing",
            "progress": 0,
            "current_step": "Initializing",
        }

        metrics = {}

        if not model_name or not hf_token or not repo_name:
            return (
                {"status": "Error", "progress": 0, "current_step": "Validation Failed"},
                "Model name, HuggingFace token, and repository name are required.",
                metrics
            )

        status["progress"] = 0.2
        status["current_step"] = "Initialization"
        status_updates.append("Initialization complete")

        quantized_model_path = None

        if quantization_type != "None":
            status.update({"progress": 0.4, "current_step": "Quantization"})
            status_updates.append(f"Applying {quantization_type} quantization")

            if not test_text:
                test_text = TASK_CONFIGS[task]["example_text"]

            try:
                handler = get_model_handler(task, model_name, quantization_type, test_text)
                quantized_model = handler.compare()
                metrics = handler.get_metrics()
                metrics = json.loads(json.dumps(metrics))

                quantized_model_path = str(resource_manager.temp_dirs["quantized"] / "model")
                quantized_model.save_pretrained(quantized_model_path)
                status_updates.append("Quantization completed successfully")
            except Exception as e:
                logger.error(f"Quantization error: {str(e)}", exc_info=True)
                return (
                    {"status": "Error", "progress": 0.4, "current_step": "Quantization Failed"},
                    f"Quantization failed: {str(e)}",
                    metrics
                )

        if enable_onnx:
            status.update({"progress": 0.6, "current_step": "ONNX Conversion"})
            status_updates.append("Converting to ONNX format")

            try:
                output_dir = str(resource_manager.temp_dirs["onnx"])
                onnx_result = convert_to_onnx(model_name, task, output_dir)

                if onnx_result is None:
                    return (
                        {"status": "Error", "progress": 0.6, "current_step": "ONNX Conversion Failed"},
                        "ONNX conversion failed.",
                        metrics
                    )

                if onnx_quantization != "None":
                    status_updates.append(f"Applying {onnx_quantization} ONNX quantization")
                    quantize_onnx_model(output_dir, onnx_quantization)

                status.update({"progress": 0.8, "current_step": "Pushing ONNX Model"})
                status_updates.append("Pushing ONNX model to Hub")
                result, push_message = push_to_hub(
                    local_path=output_dir,
                    repo_name=f"{repo_name}-optimized",
                    hf_token=hf_token,
                    tags=["onnx", "optimum", task],
                )
                status_updates.append(push_message)
            except Exception as e:
                logger.error(f"ONNX error: {str(e)}", exc_info=True)
                return (
                    {"status": "Error", "progress": 0.6, "current_step": "ONNX Processing Failed"},
                    f"ONNX processing failed: {str(e)}",
                    metrics
                )

        if quantization_type != "None" and quantized_model_path:
            status.update({"progress": 0.9, "current_step": "Pushing Quantized Model"})
            status_updates.append("Pushing quantized model to Hub")
            result, push_message = push_to_hub(
                local_path=quantized_model_path,
                repo_name=f"{repo_name}-optimized",
                hf_token=hf_token,
                tags=["quantized", task, quantization_type],
            )
            status_updates.append(push_message)

        status.update({"progress": 1.0, "status": "Complete", "current_step": "Completed"})
        cleanup_message = resource_manager.cleanup_temp_files()
        status_updates.append(cleanup_message)

        return (
            status,
            "\n".join(status_updates),
            metrics
        )

    except Exception as e:
        logger.error(f"Error during processing: {str(e)}", exc_info=True)
        return (
            {"status": "Error", "progress": 0, "current_step": "Process Failed"},
            f"An error occurred: {str(e)}",
            metrics
        )

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("""

    # 🤗 Model Conversion Hub

    Convert and optimize your Hugging Face models with quantization and ONNX support.

    """)

    with gr.Row():
        with gr.Column(scale=2):
            model_name = gr.Textbox(label="Model Name", placeholder="e.g., bert-base-uncased")
            task = gr.Dropdown(choices=list(TASK_CONFIGS.keys()), label="Task", value="text_classification")

            with gr.Group():
                gr.Markdown("### Quantization Settings")
                quantization_type = gr.Dropdown(choices=["None", "4-bit", "8-bit", "16-bit-float"], label="Quantization Type", value="None")
                test_text = gr.Textbox(label="Test Text", placeholder="Enter text for model evaluation", lines=3, visible=False)

            with gr.Group():
                gr.Markdown("### ONNX Settings")
                enable_onnx = gr.Checkbox(label="Enable ONNX Conversion")
                with gr.Group(visible=False) as onnx_group:
                    onnx_quantization = gr.Dropdown(choices=["None", "8-bit", "16-bit-int", "16-bit-float"], label="ONNX Quantization", value="None")

            with gr.Group():
                gr.Markdown("### HuggingFace Settings")
                hf_token = gr.Textbox(label="HuggingFace Token (Required)", type="password")
                repo_name = gr.Textbox(label="Repository Name")

        with gr.Column(scale=1):
            status_output = gr.JSON(label="Status", value={"status": "Ready", "progress": 0, "current_step": "Waiting"})
            message_output = gr.Markdown(label="Progress Messages")

            gr.Markdown("### Metrics")
            with gr.Group():
                metrics_output = gr.JSON(
                    value={
                        "model_sizes": {"original": 0.0, "quantized": 0.0},
                        "inference_times": {"original": 0.0, "quantized": 0.0},
                        "comparison_metrics": {}
                    },
                    show_label=True
                )

            memory_info = gr.JSON(label="Resource Usage")
            convert_btn = gr.Button("🚀 Start Conversion", variant="primary")

            with gr.Accordion("ℹ️ Help", open=False):
                gr.Markdown("""

                ### Quick Guide

                1. Enter your model name and HuggingFace token.

                2. Select the appropriate task.

                3. Choose optimization options.

                4. Click Start Conversion.



                ### Tips

                - Ensure sufficient system resources.

                - Use test text to validate conversions.

                """)

    def update_memory_info():
        resource_manager = ResourceManager()
        return resource_manager.get_memory_info()

    quantization_type.change(lambda x: gr.update(visible=x != "None"), inputs=[quantization_type], outputs=[test_text])
    task.change(lambda x: gr.update(value=TASK_CONFIGS[x]["example_text"]), inputs=[task], outputs=[test_text])
    enable_onnx.change(lambda x: gr.update(visible=x), inputs=[enable_onnx], outputs=[onnx_group])

    convert_btn.click(
        process_model,
        inputs=[model_name, task, quantization_type, enable_onnx, onnx_quantization, hf_token, repo_name, test_text],
        outputs=[status_output, message_output, metrics_output]
    )
    app.load(update_memory_info, outputs=[memory_info], every=30)

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, share=True, debug=True)