Add docstrings and reorganize the code

app.py

@@ -1,16 +1,26 @@
+"""Convert Hugging Face models to ONNX format.
+
+This application provides a Streamlit interface for converting Hugging Face models
+to ONNX format using the Transformers.js conversion scripts. It handles:
+- Model conversion with optional trust_remote_code and output_attentions
+- 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 sys
-import shutil
-import re
+from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional, Tuple
-from dataclasses import dataclass
 
 import streamlit as st
-from huggingface_hub import HfApi, whoami, model_info, hf_hub_download
 import yaml
+from huggingface_hub import HfApi, hf_hub_download, model_info, whoami
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -18,7 +28,15 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class Config:
-    """Application configuration."""
+    """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
+        repo_path: Path to the bundled transformers.js repository
+    """
 
     hf_token: str
     hf_username: str
@@ -28,10 +46,22 @@ class Config:
 
     @classmethod
     def from_env(cls) -> "Config":
-        """Create config from environment variables and secrets."""
+        """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:
@@ -39,6 +69,7 @@ class Config:
                 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:
@@ -54,53 +85,118 @@ class Config:
 
 
 class ModelConverter:
-    """Handles model conversion and upload operations."""
+    """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
+    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
+    """
 
     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:
-            path = hf_hub_download(
+            readme_path = hf_hub_download(
                 repo_id=repo_id, filename="README.md", token=self.config.hf_token
             )
-            with open(path, "r", encoding="utf-8", errors="ignore") as f:
+            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("---"):
-            m = re.match(r"^---[\s\S]*?\n---\s*\n", text)
-            if m:
-                return text[m.end() :]
+            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]:
-        """Return (frontmatter_dict, body). If no frontmatter, returns ({}, text)."""
+        """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 ""
-        m = re.match(r"^---\s*\n([\s\S]*?)\n---\s*\n", text)
-        if not m:
+
+        # 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
-        fm_text = m.group(1)
-        body = text[m.end() :]
+
+        frontmatter_text = match.group(1)
+        body = text[match.end() :]
+
+        # Parse YAML safely, returning empty dict on any error
         try:
-            data = yaml.safe_load(fm_text)
-            if not isinstance(data, dict):
-                data = {}
+            parsed_data = yaml.safe_load(frontmatter_text)
+            if not isinstance(parsed_data, dict):
+                parsed_data = {}
         except Exception:
-            data = {}
-        return data, body
+            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"
 
-    def _pipeline_docs_url(self, pipeline_tag: Optional[str]) -> Optional[str]:
-        base = "https://huggingface.co/docs/transformers.js/api/pipelines"
         if not pipeline_tag:
-            return base
-        mapping = {
+            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",
@@ -127,49 +223,85 @@ class ModelConverter:
             "image-to-image": "ImageToImagePipeline",
             "depth-estimation": "DepthEstimationPipeline",
         }
-        cls = mapping.get(pipeline_tag)
-        if not cls:
-            return base
-        return f"{base}#module_pipelines.{cls}"
 
-    def _map_pipeline_to_task(self, pipeline_tag: Optional[str]) -> Optional[str]:
+        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
-        synonyms = {
+
+        # Map abbreviated tags to their full names
+        tag_synonyms = {
             "vqa": "visual-question-answering",
         }
-        return synonyms.get(pipeline_tag, pipeline_tag)
+
+        return tag_synonyms.get(pipeline_tag, pipeline_tag)
+
+    # ============================================================================
+    # Model Conversion Methods
+    # ============================================================================
 
     def setup_repository(self) -> None:
-        """Ensure the bundled transformers.js repository is present."""
+        """Verify that the transformers.js repository exists.
+
+        Raises:
+            RuntimeError: If the repository is not found at the expected path
+        """
         if not self.config.repo_path.exists():
             raise RuntimeError(
-                f"Expected transformers.js repository at {self.config.repo_path} but it was not found."
+                f"Expected transformers.js repository at {self.config.repo_path} "
+                f"but it was not found."
             )
 
     def _run_conversion_subprocess(
-        self, input_model_id: str, extra_args: List[str] = None
+        self, input_model_id: str, extra_args: Optional[List[str]] = None
     ) -> subprocess.CompletedProcess:
-        """Run the conversion subprocess with the given arguments."""
-        cmd = [
+        """Execute the ONNX conversion script as a subprocess.
+
+        Args:
+            input_model_id: Hugging Face model ID to convert
+            extra_args: Additional command-line arguments for the conversion script
+
+        Returns:
+            subprocess.CompletedProcess: Result of the subprocess execution
+        """
+        # Build the conversion command
+        command = [
             sys.executable,
             "-m",
             "scripts.convert",
-            "--quantize",
+            "--quantize",  # Enable quantization for smaller model size
             "--model_id",
             input_model_id,
         ]
 
         if extra_args:
-            cmd.extend(extra_args)
+            command.extend(extra_args)
 
+        # Run conversion in the transformers.js repository directory
         return subprocess.run(
-            cmd,
+            command,
             cwd=self.config.repo_path,
             capture_output=True,
             text=True,
             env={
                 "HF_TOKEN": self.config.hf_token,
+                # Force eager attention implementation for compatibility
                 "TRANSFORMERS_ATTENTION_IMPLEMENTATION": "eager",
                 "PYTORCH_SDP_KERNEL": "math",
             },
@@ -178,34 +310,55 @@ class ModelConverter:
     def convert_model(
         self,
         input_model_id: str,
-        trust_remote_code=False,
-        output_attentions=False,
+        trust_remote_code: bool = False,
+        output_attentions: bool = False,
     ) -> Tuple[bool, Optional[str]]:
-        """Convert the model to ONNX format."""
+        """Convert a Hugging Face model to ONNX format.
+
+        Args:
+            input_model_id: Hugging Face model repository ID
+            trust_remote_code: Whether to trust and execute remote code from the model
+            output_attentions: Whether to output attention weights (required for some tasks)
+
+        Returns:
+            Tuple containing:
+            - bool: True if conversion succeeded, False otherwise
+            - Optional[str]: Error message if failed, or conversion log if succeeded
+        """
         try:
-            extra_args: List[str] = []
+            conversion_args: List[str] = []
+
+            # Handle trust_remote_code option (requires user token for security)
             if trust_remote_code:
                 if not self.config.is_using_user_token:
                     raise Exception(
                         "Trust Remote Code requires your own HuggingFace token."
                     )
-                extra_args.append("--trust_remote_code")
+                conversion_args.append("--trust_remote_code")
 
+            # Handle output_attentions option (needed for word-level timestamps in Whisper)
             if output_attentions:
-                extra_args.append("--output_attentions")
+                conversion_args.append("--output_attentions")
 
+            # Try to infer the task from model metadata and pass it to the conversion script
+            # This helps the script choose the right export configuration
             try:
                 info = model_info(repo_id=input_model_id, token=self.config.hf_token)
-                task = self._map_pipeline_to_task(getattr(info, "pipeline_tag", None))
+                pipeline_tag = getattr(info, "pipeline_tag", None)
+                task = self._normalize_pipeline_tag(pipeline_tag)
                 if task:
-                    extra_args.extend(["--task", task])
+                    conversion_args.extend(["--task", task])
             except Exception:
+                # If we can't fetch the task, continue without it
+                # The conversion script will try to infer it automatically
                 pass
 
+            # Run the conversion
             result = self._run_conversion_subprocess(
-                input_model_id, extra_args=extra_args or None
+                input_model_id, extra_args=conversion_args or None
             )
 
+            # Check if conversion succeeded
             if result.returncode != 0:
                 return False, result.stderr
 
@@ -214,86 +367,148 @@ class ModelConverter:
         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 model to Hugging Face."""
+        """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.config.repo_path / "models" / input_model_id
 
         try:
+            # Create the target repository (public by default)
             self.api.create_repo(output_model_id, exist_ok=True, private=False)
 
-            readme_path = f"{model_folder_path}/README.md"
-
-            with open(readme_path, "w") as file:
-                file.write(self.generate_readme(input_model_id))
+            # 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
+
+            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)
 
-    def generate_readme(self, imi: str):
+    # ============================================================================
+    # 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=imi, token=self.config.hf_token)
+            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
 
-        original_text = self._fetch_original_readme(imi)
+        # 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"] = [imi]
+        merged_meta["base_model"] = [input_model_id]
         if pipeline_tag is not None:
             merged_meta["pipeline_tag"] = pipeline_tag
 
-        fm_yaml = yaml.safe_dump(merged_meta, sort_keys=False).strip()
-        header = f"---\n{fm_yaml}\n---\n\n"
+        # Generate YAML frontmatter
+        frontmatter_yaml = yaml.safe_dump(merged_meta, sort_keys=False).strip()
+        header = f"---\n{frontmatter_yaml}\n---\n\n"
 
-        parts: List[str] = []
-        parts.append(header)
-        parts.append(f"# {imi.split('/')[-1]} (ONNX)\n")
-        parts.append(
-            f"This is an ONNX version of [{imi}](https://huggingface.co/{imi}). "
+        # 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)."
         )
 
-        docs_url = self._pipeline_docs_url(pipeline_tag)
+        # Add usage section with Transformers.js docs link
+        docs_url = self._get_pipeline_docs_url(pipeline_tag)
         if docs_url:
-            parts.append("\n## Usage with Transformers.js\n")
+            readme_sections.append("\n## Usage with Transformers.js\n")
             if pipeline_tag:
-                parts.append(
+                readme_sections.append(
                     f"See the pipeline documentation for `{pipeline_tag}`: {docs_url}"
                 )
             else:
-                parts.append(f"See the pipelines documentation: {docs_url}")
+                readme_sections.append(f"See the pipelines documentation: {docs_url}")
 
+        # Append original README content (if available)
         if original_body:
-            parts.append("\n---\n")
-            parts.append(original_body)
+            readme_sections.append("\n---\n")
+            readme_sections.append(original_body)
 
-        return "\n\n".join(parts) + "\n"
+        return "\n\n".join(readme_sections) + "\n"
 
 
 def main():
-    """Main application entry point."""
+    """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)
         converter.setup_repository()
 
+        # Get model ID from user
         input_model_id = st.text_input(
             "Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"
         )
@@ -301,23 +516,29 @@ def main():
         if not input_model_id:
             return
 
+        # Optional: User token input
         st.text_input(
             f"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: Output attentions (for Whisper models)
         output_attentions = False
         if "whisper" in input_model_id.lower():
             output_attentions = st.toggle(
                 "Whether to output attentions from the Whisper model. This is required for word-level (token) timestamps."
             )
 
+        # 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"
@@ -326,25 +547,29 @@ def main():
             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(f"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,
@@ -358,6 +583,7 @@ def main():
             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: