app.py

"""Convert Hugging Face models to ONNX format.

This application provides a Streamlit interface for converting Hugging Face models
to ONNX format using Optimum's ONNX exporter. It handles:
- Model conversion with optional trust_remote_code
- Automatic task inference with fallback support
- README generation with merged metadata from the original model
- Upload to Hugging Face Hub
"""

import logging
import os
import re
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

import onnx
from onnxconverter_common import float16 as onnx_float16
from onnxruntime.quantization import QuantType, quantize_dynamic
from onnxruntime.quantization.matmul_bnb4_quantizer import MatMulBnb4Quantizer
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer

import streamlit as st
import yaml
from huggingface_hub import HfApi, hf_hub_download, model_info, whoami

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class Config:
    """Application configuration containing authentication and path settings.

    Attributes:
        hf_token: Hugging Face API token (user token takes precedence over system token)
        hf_username: Hugging Face username associated with the token
        is_using_user_token: True if using a user-provided token, False if using system token
        hf_base_url: Base URL for Hugging Face Hub
    """

    hf_token: str
    hf_username: str
    is_using_user_token: bool
    hf_base_url: str = "https://huggingface.co"

    @classmethod
    def from_env(cls) -> "Config":
        """Create configuration from environment variables and Streamlit session state.

        Priority order for tokens:
        1. User-provided token from Streamlit session (st.session_state.user_hf_token)
        2. System token from environment variable (HF_TOKEN)

        Returns:
            Config: Initialized configuration object

        Raises:
            ValueError: If no valid token is available
        """
        system_token = os.getenv("HF_TOKEN")
        user_token = st.session_state.get("user_hf_token")

        # Determine username based on which token is being used
        if user_token:
            hf_username = whoami(token=user_token)["name"]
        else:
            hf_username = (
                os.getenv("SPACE_AUTHOR_NAME") or whoami(token=system_token)["name"]
            )

        # User token takes precedence over system token
        hf_token = user_token or system_token

        if not hf_token:
            raise ValueError(
                "When the user token is not provided, the system token must be set."
            )

        return cls(
            hf_token=hf_token,
            hf_username=hf_username,
            is_using_user_token=bool(user_token),
        )


class ModelConverter:
    """Handles model conversion to ONNX format and upload to Hugging Face Hub.

    This class manages the entire conversion workflow:
    1. Fetching original model metadata and README
    2. Running the ONNX conversion subprocess via optimum-cli
    3. Generating an enhanced README with merged metadata
    4. Uploading the converted model to Hugging Face Hub

    Attributes:
        config: Application configuration containing tokens and paths
        api: Hugging Face API client for repository operations
    """

    OUTPUT_BASE = Path("./onnx_output")

    def __init__(self, config: Config):
        """Initialize the converter with configuration.

        Args:
            config: Application configuration object
        """
        self.config = config
        self.api = HfApi(token=config.hf_token)

    # ============================================================================
    # README Processing Methods
    # ============================================================================

    def _fetch_original_readme(self, repo_id: str) -> str:
        """Download the README from the original model repository.

        Args:
            repo_id: Hugging Face model repository ID (e.g., 'username/model-name')

        Returns:
            str: Content of the README file, or empty string if not found
        """
        try:
            readme_path = hf_hub_download(
                repo_id=repo_id, filename="README.md", token=self.config.hf_token
            )
            with open(readme_path, "r", encoding="utf-8", errors="ignore") as f:
                return f.read()
        except Exception:
            # Silently fail if README doesn't exist or can't be downloaded
            return ""

    def _strip_yaml_frontmatter(self, text: str) -> str:
        """Remove YAML frontmatter from text, returning only the body.

        YAML frontmatter is delimited by '---' at the start and end.

        Args:
            text: Text that may contain YAML frontmatter

        Returns:
            str: Text with frontmatter removed, or original text if no frontmatter found
        """
        if not text:
            return ""
        if text.startswith("---"):
            match = re.match(r"^---[\s\S]*?\n---\s*\n", text)
            if match:
                return text[match.end() :]
        return text

    def _extract_yaml_frontmatter(self, text: str) -> Tuple[dict, str]:
        """Parse and extract YAML frontmatter from text.

        Args:
            text: Text that may contain YAML frontmatter

        Returns:
            Tuple containing:
            - dict: Parsed YAML frontmatter as a dictionary (empty dict if none found)
            - str: Remaining body text after the frontmatter
        """
        if not text or not text.startswith("---"):
            return {}, text or ""

        # Match YAML frontmatter pattern: ---\n...content...\n---\n
        match = re.match(r"^---\s*\n([\s\S]*?)\n---\s*\n", text)
        if not match:
            return {}, text

        frontmatter_text = match.group(1)
        body = text[match.end() :]

        # Parse YAML safely, returning empty dict on any error
        try:
            parsed_data = yaml.safe_load(frontmatter_text)
            if not isinstance(parsed_data, dict):
                parsed_data = {}
        except Exception:
            parsed_data = {}

        return parsed_data, body

    def _get_pipeline_docs_url(self, pipeline_tag: Optional[str]) -> str:
        """Generate Transformers.js documentation URL for a given pipeline tag.

        Args:
            pipeline_tag: Hugging Face pipeline tag (e.g., 'text-generation')

        Returns:
            str: URL to the relevant Transformers.js pipeline documentation
        """
        base_url = "https://huggingface.co/docs/transformers.js/api/pipelines"

        if not pipeline_tag:
            return base_url

        # Map Hugging Face pipeline tags to Transformers.js pipeline class names
        pipeline_class_mapping = {
            "text-classification": "TextClassificationPipeline",
            "token-classification": "TokenClassificationPipeline",
            "question-answering": "QuestionAnsweringPipeline",
            "fill-mask": "FillMaskPipeline",
            "text2text-generation": "Text2TextGenerationPipeline",
            "summarization": "SummarizationPipeline",
            "translation": "TranslationPipeline",
            "text-generation": "TextGenerationPipeline",
            "zero-shot-classification": "ZeroShotClassificationPipeline",
            "feature-extraction": "FeatureExtractionPipeline",
            "image-feature-extraction": "ImageFeatureExtractionPipeline",
            "audio-classification": "AudioClassificationPipeline",
            "zero-shot-audio-classification": "ZeroShotAudioClassificationPipeline",
            "automatic-speech-recognition": "AutomaticSpeechRecognitionPipeline",
            "image-to-text": "ImageToTextPipeline",
            "image-classification": "ImageClassificationPipeline",
            "image-segmentation": "ImageSegmentationPipeline",
            "background-removal": "BackgroundRemovalPipeline",
            "zero-shot-image-classification": "ZeroShotImageClassificationPipeline",
            "object-detection": "ObjectDetectionPipeline",
            "zero-shot-object-detection": "ZeroShotObjectDetectionPipeline",
            "document-question-answering": "DocumentQuestionAnsweringPipeline",
            "text-to-audio": "TextToAudioPipeline",
            "image-to-image": "ImageToImagePipeline",
            "depth-estimation": "DepthEstimationPipeline",
        }

        pipeline_class = pipeline_class_mapping.get(pipeline_tag)
        if not pipeline_class:
            return base_url

        return f"{base_url}#module_pipelines.{pipeline_class}"

    def _normalize_pipeline_tag(self, pipeline_tag: Optional[str]) -> Optional[str]:
        """Normalize pipeline tag to match expected task names.

        Some pipeline tags use abbreviations that need to be expanded
        for the conversion script to recognize them.

        Args:
            pipeline_tag: Original pipeline tag from model metadata

        Returns:
            Optional[str]: Normalized task name, or None if input is None
        """
        if not pipeline_tag:
            return None

        # Map abbreviated tags to their full names
        tag_synonyms = {
            "vqa": "visual-question-answering",
        }

        return tag_synonyms.get(pipeline_tag, pipeline_tag)

    # ============================================================================
    # Model Conversion Methods
    # ============================================================================

    def _run_export_subprocess(
        self,
        input_model_id: str,
        output_path: Path,
        extra_args: Optional[List[str]] = None,
    ) -> subprocess.CompletedProcess:
        """Execute the optimum-cli ONNX export as a subprocess.

        Args:
            input_model_id: Hugging Face model ID to convert
            output_path: Directory where the exported files will be written
            extra_args: Additional command-line arguments for optimum-cli

        Returns:
            subprocess.CompletedProcess: Result of the subprocess execution
        """
        command = [
            "optimum-cli", "export", "onnx",
            "--model", input_model_id,
            str(output_path),
        ]

        if extra_args:
            command.extend(extra_args)

        env = os.environ.copy()
        env["HF_TOKEN"] = self.config.hf_token

        return subprocess.run(
            command,
            capture_output=True,
            text=True,
            env=env,
        )

    def _export_base(
        self,
        input_model_id: str,
        output_path: Path,
        extra_args: Optional[List[str]] = None,
    ) -> Tuple[bool, str]:
        """Export the fp32 model, placing ONNX files in onnx/ and config files in root.

        Args:
            input_model_id: Hugging Face model ID to convert
            output_path: Root output directory for the converted repository
            extra_args: Additional command-line arguments for optimum-cli

        Returns:
            Tuple of (success, log/error message)
        """
        with tempfile.TemporaryDirectory() as tmp:
            tmp_path = Path(tmp)
            result = self._run_export_subprocess(input_model_id, tmp_path, extra_args=extra_args)
            if result.returncode != 0:
                return False, result.stderr

            onnx_dir = output_path / "onnx"
            onnx_dir.mkdir(parents=True, exist_ok=True)

            for file in tmp_path.iterdir():
                if file.suffix in (".onnx", ".onnx_data"):
                    shutil.copy2(file, onnx_dir / file.name)
                else:
                    shutil.copy2(file, output_path / file.name)

            return True, result.stderr

    def _apply_quantizations(self, onnx_dir: Path) -> str:
        """Apply all quantization variants to each base ONNX file in onnx_dir.

        Produces the following variants for each base file (e.g. model.onnx):
            model_fp16.onnx, model_int8.onnx, model_uint8.onnx,
            model_quantized.onnx, model_q4.onnx, model_q4f16.onnx, model_bnb4.onnx

        Each variant failure is non-fatal; the log records per-variant outcomes.

        Args:
            onnx_dir: Directory containing the base ONNX file(s)

        Returns:
            str: Log of outcomes for each quantization variant
        """
        _VARIANT_SUFFIXES = ("_fp16", "_int8", "_uint8", "_quantized", "_q4", "_q4f16", "_bnb4")

        logs: List[str] = []

        for base_file in sorted(onnx_dir.glob("*.onnx")):
            if any(base_file.stem.endswith(s) for s in _VARIANT_SUFFIXES):
                continue
            stem = base_file.stem
            model = onnx.load(str(base_file))

            # fp16: convert float32 weights to float16
            try:
                fp16_model = onnx_float16.convert_float_to_float16(model, keep_io_types=True)
                onnx.save(fp16_model, str(onnx_dir / f"{stem}_fp16.onnx"))
                logs.append(f"{stem}_fp16: ok")
            except Exception as e:
                logs.append(f"{stem}_fp16: failed ({e})")

            # int8: dynamic quantization with signed weights
            try:
                quantize_dynamic(
                    model_input=str(base_file),
                    model_output=str(onnx_dir / f"{stem}_int8.onnx"),
                    weight_type=QuantType.QInt8,
                )
                logs.append(f"{stem}_int8: ok")
            except Exception as e:
                logs.append(f"{stem}_int8: failed ({e})")

            # uint8: dynamic quantization with unsigned weights
            try:
                quantize_dynamic(
                    model_input=str(base_file),
                    model_output=str(onnx_dir / f"{stem}_uint8.onnx"),
                    weight_type=QuantType.QUInt8,
                )
                logs.append(f"{stem}_uint8: ok")
            except Exception as e:
                logs.append(f"{stem}_uint8: failed ({e})")

            # quantized: alias for int8 (fallback to uint8)
            int8_path = onnx_dir / f"{stem}_int8.onnx"
            uint8_path = onnx_dir / f"{stem}_uint8.onnx"
            quantized_src = int8_path if int8_path.exists() else (uint8_path if uint8_path.exists() else None)
            if quantized_src:
                shutil.copy2(quantized_src, onnx_dir / f"{stem}_quantized.onnx")
                logs.append(f"{stem}_quantized: ok (copy of {quantized_src.name})")

            # q4 and q4f16: 4-bit symmetric quantization (nbits=4)
            try:
                q4_quantizer = MatMulNBitsQuantizer(model, bits=4, block_size=32, is_symmetric=True)
                q4_quantizer.process()
                q4_model = q4_quantizer.model.model
                q4_path = onnx_dir / f"{stem}_q4.onnx"
                onnx.save(q4_model, str(q4_path))
                logs.append(f"{stem}_q4: ok")

                try:
                    q4f16_model = onnx_float16.convert_float_to_float16(q4_model, keep_io_types=True)
                    onnx.save(q4f16_model, str(onnx_dir / f"{stem}_q4f16.onnx"))
                    logs.append(f"{stem}_q4f16: ok")
                except Exception as e:
                    logs.append(f"{stem}_q4f16: failed ({e})")
            except Exception as e:
                logs.append(f"{stem}_q4: failed ({e})")
                logs.append(f"{stem}_q4f16: skipped (q4 failed)")

            # bnb4: BitsAndBytes NF4 quantization
            try:
                bnb4_quantizer = MatMulBnb4Quantizer(
                    model, block_size=64, quant_type=MatMulBnb4Quantizer.NF4
                )
                bnb4_quantizer.process()
                onnx.save(bnb4_quantizer.model.model, str(onnx_dir / f"{stem}_bnb4.onnx"))
                logs.append(f"{stem}_bnb4: ok")
            except Exception as e:
                logs.append(f"{stem}_bnb4: failed ({e})")

        return "\n".join(logs)

    def convert_model(
        self,
        input_model_id: str,
        trust_remote_code: bool = False,
        enable_task_inference: bool = True,
    ) -> Tuple[bool, Optional[str]]:
        """Convert a Hugging Face model to ONNX format using optimum-cli.

        Produces multiple quantization variants, organized as:
            <model>/
                config.json, tokenizer.json, ...   (config files)
                onnx/
                    model.onnx                     (fp32)
                    model_fp16.onnx
                    model_int8.onnx
                    model_uint8.onnx
                    model_quantized.onnx           (copy of int8)
                    model_q4.onnx
                    model_q4f16.onnx
                    model_bnb4.onnx

        Args:
            input_model_id: Hugging Face model repository ID
            trust_remote_code: Whether to trust and execute remote code from the model
            enable_task_inference: Whether to pass the task argument based on the model's pipeline tag

        Returns:
            Tuple containing:
            - bool: True if conversion succeeded, False otherwise
            - Optional[str]: Error message if failed, or conversion log if succeeded
        """
        try:
            base_args: List[str] = []

            if trust_remote_code:
                if not self.config.is_using_user_token:
                    raise Exception(
                        "Trust Remote Code requires your own HuggingFace token."
                    )
                base_args.append("--trust-remote-code")

            if enable_task_inference:
                try:
                    info = model_info(
                        repo_id=input_model_id, token=self.config.hf_token
                    )
                    pipeline_tag = getattr(info, "pipeline_tag", None)
                    task = self._normalize_pipeline_tag(pipeline_tag)
                    if task:
                        base_args.extend(["--task", task])
                except Exception:
                    pass

            output_path = self.OUTPUT_BASE / input_model_id
            output_path.mkdir(parents=True, exist_ok=True)

            # Export fp32 base model; organizes output into root + onnx/ subfolder
            success, log = self._export_base(
                input_model_id, output_path, extra_args=base_args
            )
            if not success:
                return False, log

            # Apply all quantization variants to each exported ONNX file
            quant_log = self._apply_quantizations(output_path / "onnx")

            return True, log + "\n\nQuantization results:\n" + quant_log

        except Exception as e:
            return False, str(e)

    # ============================================================================
    # Upload Methods
    # ============================================================================

    def upload_model(self, input_model_id: str, output_model_id: str) -> Optional[str]:
        """Upload the converted ONNX model to Hugging Face Hub.

        This method:
        1. Creates the target repository (if it doesn't exist)
        2. Generates an enhanced README with merged metadata
        3. Uploads all model files to the repository
        4. Cleans up local files after upload

        Args:
            input_model_id: Original model repository ID
            output_model_id: Target repository ID for the ONNX model

        Returns:
            Optional[str]: Error message if upload failed, None if successful
        """
        model_folder_path = self.OUTPUT_BASE / input_model_id

        try:
            # Create the target repository (public by default)
            self.api.create_repo(output_model_id, exist_ok=True, private=False)

            # Generate and write the enhanced README
            readme_path = model_folder_path / "README.md"
            readme_content = self.generate_readme(input_model_id)
            readme_path.write_text(readme_content, encoding="utf-8")

            # Upload all files from the model folder
            self.api.upload_folder(
                folder_path=str(model_folder_path), repo_id=output_model_id
            )

            return None  # Success

        except Exception as e:
            return str(e)
        finally:
            # Always clean up local files, even if upload failed
            shutil.rmtree(model_folder_path, ignore_errors=True)

    # ============================================================================
    # README Generation Methods
    # ============================================================================

    def generate_readme(self, input_model_id: str) -> str:
        """Generate an enhanced README for the ONNX model.

        This method creates a README that:
        1. Merges metadata from the original model with ONNX-specific metadata
        2. Adds a description and link to the conversion space
        3. Includes usage instructions with links to Transformers.js docs
        4. Appends the original model's README content

        Args:
            input_model_id: Original model repository ID

        Returns:
            str: Complete README content in Markdown format with YAML frontmatter
        """
        # Fetch pipeline tag from model metadata (if available)
        try:
            info = model_info(repo_id=input_model_id, token=self.config.hf_token)
            pipeline_tag = getattr(info, "pipeline_tag", None)
        except Exception:
            pipeline_tag = None

        # Fetch and parse the original README
        original_text = self._fetch_original_readme(input_model_id)
        original_meta, original_body = self._extract_yaml_frontmatter(original_text)
        original_body = (
            original_body or self._strip_yaml_frontmatter(original_text)
        ).strip()

        # Merge original metadata with our ONNX-specific metadata (ours take precedence)
        merged_meta = {}
        if isinstance(original_meta, dict):
            merged_meta.update(original_meta)
        merged_meta["library_name"] = "transformers.js"
        merged_meta["base_model"] = [input_model_id]
        if pipeline_tag is not None:
            merged_meta["pipeline_tag"] = pipeline_tag

        # Generate YAML frontmatter
        frontmatter_yaml = yaml.safe_dump(merged_meta, sort_keys=False).strip()
        header = f"---\n{frontmatter_yaml}\n---\n\n"

        # Build README sections
        readme_sections: List[str] = []
        readme_sections.append(header)

        # Add title
        model_name = input_model_id.split("/")[-1]
        readme_sections.append(f"# {model_name} (ONNX)\n")

        # Add description
        readme_sections.append(
            f"This is an ONNX version of [{input_model_id}](https://huggingface.co/{input_model_id}). "
            "It was automatically converted and uploaded using "
            "[this Hugging Face Space](https://huggingface.co/spaces/onnx-community/convert-to-onnx)."
        )

        # Add usage section with Transformers.js docs link
        docs_url = self._get_pipeline_docs_url(pipeline_tag)
        if docs_url:
            readme_sections.append("\n## Usage with Transformers.js\n")
            if pipeline_tag:
                readme_sections.append(
                    f"See the pipeline documentation for `{pipeline_tag}`: {docs_url}"
                )
            else:
                readme_sections.append(f"See the pipelines documentation: {docs_url}")

        # Append original README content (if available)
        if original_body:
            readme_sections.append("\n---\n")
            readme_sections.append(original_body)

        return "\n\n".join(readme_sections) + "\n"


def main():
    """Main application entry point for the Streamlit interface.

    This function:
    1. Initializes configuration and converter
    2. Displays the UI for model input and options
    3. Handles the conversion workflow
    4. Shows progress and results to the user
    """
    st.write("## Convert a Hugging Face model to ONNX")

    try:
        # Initialize configuration and converter
        config = Config.from_env()
        converter = ModelConverter(config)

        # Get model ID from user
        input_model_id = st.text_input(
            "Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"
        )

        if not input_model_id:
            return

        # Optional: User token input
        st.text_input(
            "Optional: Your Hugging Face write token. Fill it if you want to upload the model under your account.",
            type="password",
            key="user_hf_token",
        )

        # Optional: Trust remote code toggle (requires user token)
        trust_remote_code = st.toggle("Optional: Trust Remote Code.")
        if trust_remote_code:
            st.warning(
                "This option should only be enabled for repositories you trust and in which you have read the code, as it will execute arbitrary code present in the model repository. When this option is enabled, you must use your own Hugging Face write token."
            )

        # Optional: Task inference toggle
        enable_task_inference = st.toggle(
            "Optional: Base the 'task' argument from the conversion script on the model's pipeline tag",
            value=False,
            help="This can make the conversion of some models work, but may cause issues for others. It's recommended to first try converting the model with this option disabled, and only enable it if the conversion fails.",
        )

        # Determine output repository
        # If user owns the model, allow uploading to the same repo
        if config.hf_username == input_model_id.split("/")[0]:
            same_repo = st.checkbox(
                "Upload the ONNX weights to the existing repository"
            )
        else:
            same_repo = False

        model_name = input_model_id.split("/")[-1]
        output_model_id = f"{config.hf_username}/{model_name}"

        # Add -ONNX suffix if creating a new repository
        if not same_repo:
            output_model_id += "-ONNX"

        output_model_url = f"{config.hf_base_url}/{output_model_id}"

        # Check if model already exists
        if not same_repo and converter.api.repo_exists(output_model_id):
            st.write("This model has already been converted! 🎉")
            st.link_button(f"Go to {output_model_id}", output_model_url, type="primary")
            return

        # Show where the model will be uploaded
        st.write("URL where the model will be converted and uploaded to:")
        st.code(output_model_url, language="plaintext")

        # Wait for user confirmation before proceeding
        if not st.button(label="Proceed", type="primary"):
            return

        # Step 1: Convert the model to ONNX
        with st.spinner("Converting model..."):
            success, stderr = converter.convert_model(
                input_model_id,
                trust_remote_code=trust_remote_code,
                enable_task_inference=enable_task_inference,
            )
            if not success:
                st.error(f"Conversion failed: {stderr}")
                return

            st.success("Conversion successful!")
            st.code(stderr)

        # Step 2: Upload the converted model to Hugging Face
        with st.spinner("Uploading model..."):
            error = converter.upload_model(input_model_id, output_model_id)
            if error:
                st.error(f"Upload failed: {error}")
                return

            st.success("Upload successful!")
            st.write("You can now go and view the model on Hugging Face!")
            st.link_button(f"Go to {output_model_id}", output_model_url, type="primary")

    except Exception as e:
        logger.exception("Application error")
        st.error(f"An error occurred: {str(e)}")


if __name__ == "__main__":
    main()