"""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: / 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()