Refactor the code, making it more robust and maintainable

app.py

@@ -1,139 +1,211 @@
+import logging
 import os
 import subprocess
 import sys
-import tarfile
-import tempfile
-import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Tuple
+from urllib.request import urlopen, urlretrieve
 
 import streamlit as st
 from huggingface_hub import HfApi
 
-HF_TOKEN = st.secrets.get("HF_TOKEN") or os.environ.get("HF_TOKEN")
-HF_USERNAME = (
-    st.secrets.get("HF_USERNAME")
-    or os.environ.get("HF_USERNAME")
-    or os.environ.get("SPACE_AUTHOR_NAME")
-)
-
-TRANSFORMERS_BASE_URL = "https://github.com/xenova/transformers.js/archive/refs"
-TRANSFORMERS_REPOSITORY_REVISION = "3.0.0"
-TRANSFORMERS_REF_TYPE = (
-    "tags"
-    if urllib.request.urlopen(
-        f"{TRANSFORMERS_BASE_URL}/tags/{TRANSFORMERS_REPOSITORY_REVISION}.tar.gz"
-    ).getcode()
-    == 200
-    else "heads"
-)
-TRANSFORMERS_REPOSITORY_URL = f"{TRANSFORMERS_BASE_URL}/{TRANSFORMERS_REF_TYPE}/{TRANSFORMERS_REPOSITORY_REVISION}.tar.gz"
-TRANSFORMERS_REPOSITORY_PATH = "./transformers.js"
-ARCHIVE_PATH = f"./transformers_{TRANSFORMERS_REPOSITORY_REVISION}.tar.gz"
-HF_BASE_URL = "https://huggingface.co"
-
-if not os.path.exists(TRANSFORMERS_REPOSITORY_PATH):
-    urllib.request.urlretrieve(TRANSFORMERS_REPOSITORY_URL, ARCHIVE_PATH)
-
-    with tempfile.TemporaryDirectory() as tmp_dir:
-        with tarfile.open(ARCHIVE_PATH, "r:gz") as tar:
-            tar.extractall(tmp_dir)
-
-        extracted_folder = os.path.join(tmp_dir, os.listdir(tmp_dir)[0])
-
-        os.rename(extracted_folder, TRANSFORMERS_REPOSITORY_PATH)
-
-    os.remove(ARCHIVE_PATH)
-    print("Repository downloaded and extracted successfully.")
-
-st.write("## Convert a HuggingFace model to ONNX")
-
-input_model_id = st.text_input(
-    "Enter the HuggingFace model ID to convert. Example: `EleutherAI/pythia-14m`"
-)
-
-if input_model_id:
-    model_name = (
-        input_model_id.replace(f"{HF_BASE_URL}/", "")
-        .replace("/", "-")
-        .replace(f"{HF_USERNAME}-", "")
-        .strip()
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Config:
+    """Application configuration."""
+
+    hf_token: str
+    hf_username: str
+    transformers_version: str = "3.0.0"
+    hf_base_url: str = "https://huggingface.co"
+    transformers_base_url: str = (
+        "https://github.com/xenova/transformers.js/archive/refs"
     )
-    output_model_id = f"{HF_USERNAME}/{model_name}-ONNX"
-    output_model_url = f"{HF_BASE_URL}/{output_model_id}"
-    api = HfApi(token=HF_TOKEN)
-    repo_exists = api.repo_exists(output_model_id)
-
-    if repo_exists:
-        st.write("This model has already been converted! 🎉")
-        st.link_button(f"Go to {output_model_id}", output_model_url, type="primary")
-    else:
-        st.write(f"This model will be converted and uploaded to the following URL:")
-        st.code(output_model_url, language="plaintext")
-        start_conversion = st.button(label="Proceed", type="primary")
-
-        if start_conversion:
-            with st.spinner("Converting model..."):
-                output = subprocess.run(
-                    [
-                        sys.executable,
-                        "-m",
-                        "scripts.convert",
-                        "--quantize",
-                        "--model_id",
-                        input_model_id,
-                    ],
-                    cwd=TRANSFORMERS_REPOSITORY_PATH,
-                    capture_output=True,
-                    text=True,
-                    env={},
-                )
-
-                # Log the script output
-                print("### Script Output ###")
-                print(output.stdout)
-
-                # Log any errors
-                if output.stderr:
-                    print("### Script Errors ###")
-                    print(output.stderr)
-
-            model_folder_path = (
-                f"{TRANSFORMERS_REPOSITORY_PATH}/models/{input_model_id}"
+    repo_path: Path = Path("./transformers.js")
+
+    @classmethod
+    def from_env(cls) -> "Config":
+        """Create config from environment variables and secrets."""
+        hf_token = st.secrets.get("HF_TOKEN") or os.getenv("HF_TOKEN", "")
+        hf_username = (
+            st.secrets.get("HF_USERNAME")
+            or os.getenv("HF_USERNAME")
+            or os.getenv("SPACE_AUTHOR_NAME", "")
+        )
+
+        if not hf_token or not hf_username:
+            raise ValueError("HF_TOKEN and HF_USERNAME must be set")
+
+        return cls(hf_token=hf_token, hf_username=hf_username)
+
+
+class ModelConverter:
+    """Handles model conversion and upload operations."""
+
+    def __init__(self, config: Config):
+        self.config = config
+        self.api = HfApi(token=config.hf_token)
+
+    def _get_ref_type(self) -> str:
+        """Determine the reference type for the transformers repository."""
+        url = f"{self.config.transformers_base_url}/tags/{self.config.transformers_version}.tar.gz"
+        try:
+            return "tags" if urlopen(url).getcode() == 200 else "heads"
+        except Exception as e:
+            logger.warning(f"Failed to check tags, defaulting to heads: {e}")
+            return "heads"
+
+    def setup_repository(self) -> None:
+        """Download and setup transformers repository if needed."""
+        if self.config.repo_path.exists():
+            return
+
+        ref_type = self._get_ref_type()
+        archive_url = f"{self.config.transformers_base_url}/{ref_type}/{self.config.transformers_version}.tar.gz"
+        archive_path = Path(f"./transformers_{self.config.transformers_version}.tar.gz")
+
+        try:
+            urlretrieve(archive_url, archive_path)
+            self._extract_archive(archive_path)
+            logger.info("Repository downloaded and extracted successfully")
+        except Exception as e:
+            raise RuntimeError(f"Failed to setup repository: {e}")
+        finally:
+            archive_path.unlink(missing_ok=True)
+
+    def _extract_archive(self, archive_path: Path) -> None:
+        """Extract the downloaded archive."""
+        import tarfile
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            with tarfile.open(archive_path, "r:gz") as tar:
+                tar.extractall(tmp_dir)
+
+            extracted_folder = next(Path(tmp_dir).iterdir())
+            extracted_folder.rename(self.config.repo_path)
+
+    def convert_model(self, input_model_id: str) -> Tuple[bool, Optional[str]]:
+        """Convert the model to ONNX format."""
+        try:
+            result = subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "scripts.convert",
+                    "--quantize",
+                    "--model_id",
+                    input_model_id,
+                ],
+                cwd=self.config.repo_path,
+                capture_output=True,
+                text=True,
+                env={},
             )
 
-            os.rename(
-                f"{model_folder_path}/onnx/model.onnx",
-                f"{model_folder_path}/onnx/decoder_model_merged.onnx",
-            )
-            os.rename(
-                f"{model_folder_path}/onnx/model_quantized.onnx",
-                f"{model_folder_path}/onnx/decoder_model_merged_quantized.onnx",
+            if result.returncode != 0:
+                return False, result.stderr
+
+            self._rename_model_files(input_model_id)
+            return True, result.stderr
+
+        except Exception as e:
+            return False, str(e)
+
+    def _rename_model_files(self, input_model_id: str) -> None:
+        """Rename the converted model files."""
+        model_path = self.config.repo_path / "models" / input_model_id / "onnx"
+
+        renames = [
+            ("model.onnx", "decoder_model_merged.onnx"),
+            ("model_quantized.onnx", "decoder_model_merged_quantized.onnx"),
+        ]
+
+        for old_name, new_name in renames:
+            (model_path / old_name).rename(model_path / new_name)
+
+    def upload_model(self, input_model_id: str, output_model_id: str) -> Optional[str]:
+        """Upload the converted model to Hugging Face."""
+        try:
+            self.api.create_repo(output_model_id, exist_ok=True, private=False)
+            model_folder_path = self.config.repo_path / "models" / input_model_id
+
+            self.api.upload_folder(
+                folder_path=str(model_folder_path), repo_id=output_model_id
             )
+            return None
+        except Exception as e:
+            return str(e)
+        finally:
+            import shutil
 
-            st.success("Conversion successful!")
+            shutil.rmtree(model_folder_path, ignore_errors=True)
 
-            st.code(output.stderr)
 
-            with st.spinner("Uploading model..."):
-                repository = api.create_repo(
-                    f"{output_model_id}", exist_ok=True, private=False
-                )
+def main():
+    """Main application entry point."""
+    st.write("## Convert a Hugging Face model to ONNX")
+
+    try:
+        config = Config.from_env()
+        converter = ModelConverter(config)
+        converter.setup_repository()
+
+        input_model_id = st.text_input(
+            "Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"
+        )
+
+        if not input_model_id:
+            return
+
+        model_name = (
+            input_model_id.replace(f"{config.hf_base_url}/", "")
+            .replace("/", "-")
+            .replace(f"{config.hf_username}-", "")
+            .strip()
+        )
+
+        output_model_id = f"{config.hf_username}/{model_name}-ONNX"
+        output_model_url = f"{config.hf_base_url}/{output_model_id}"
+
+        if 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
+
+        st.write(f"This model will be converted and uploaded to the following URL:")
+        st.code(output_model_url, language="plaintext")
+
+        if not st.button(label="Proceed", type="primary"):
+            return
+
+        with st.spinner("Converting model..."):
+            success, stderr = converter.convert_model(input_model_id)
+            if not success:
+                st.error(f"Conversion failed: {stderr}")
+                return
+
+            st.success("Conversion successful!")
+            st.code(stderr)
+
+        with st.spinner("Uploading model..."):
+            error = converter.upload_model(input_model_id, output_model_id)
+            if error:
+                st.error(f"Upload failed: {error}")
+                return
 
-                upload_error_message = None
+            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")
 
-                try:
-                    api.upload_folder(
-                        folder_path=model_folder_path, repo_id=repository.repo_id
-                    )
-                except Exception as e:
-                    upload_error_message = str(e)
+    except Exception as e:
+        logger.exception("Application error")
+        st.error(f"An error occurred: {str(e)}")
 
-            os.system(f"rm -rf {model_folder_path}")
 
-            if upload_error_message:
-                st.error(f"Upload failed: {upload_error_message}")
-            else:
-                st.success(f"Upload successful!")
-                st.write("You can now go and view the model on HuggingFace!")
-                st.link_button(
-                    f"Go to {output_model_id}", output_model_url, type="primary"
-                )
+if __name__ == "__main__":
+    main()