Update app.py

app.py

@@ -2,7 +2,6 @@ import logging
 import os
 import subprocess
 import sys
-import warnings
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, Tuple
@@ -10,11 +9,6 @@ from urllib.request import urlopen, urlretrieve
 
 import streamlit as st
 from huggingface_hub import HfApi, whoami
-from torch.jit import TracerWarning
-from transformers import AutoConfig, GenerationConfig
-
-# Suppress local TorchScript tracer warnings
-warnings.filterwarnings("ignore", category=TracerWarning)
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -23,6 +17,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class Config:
     """Application configuration."""
+
     hf_token: str
     hf_username: str
     transformers_version: str = "3.5.0"
@@ -44,8 +39,10 @@ class Config:
                 os.getenv("SPACE_AUTHOR_NAME") or whoami(token=system_token)["name"]
             )
         hf_token = user_token or system_token
+
         if not hf_token:
             raise ValueError("HF_TOKEN must be set")
+
         return cls(hf_token=hf_token, hf_username=hf_username)
 
 
@@ -66,12 +63,14 @@ class ModelConverter:
             return "heads"
 
     def setup_repository(self) -> None:
-        """Download and setup transformers.js repo if needed."""
+        """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)
@@ -83,38 +82,19 @@ class ModelConverter:
 
     def _extract_archive(self, archive_path: Path) -> None:
         """Extract the downloaded archive."""
-        import tarfile, tempfile
+        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, always exporting attention maps.
-        Relocate generation params, suppress tracer warnings, and
-        filter out relocation/tracer warnings from stderr.
-        """
+        """Convert the model to ONNX format, always exporting attention maps."""
         try:
-            # 1. Prepare a local folder for config tweaks
-            model_dir = self.config.repo_path / "models" / input_model_id
-            model_dir.mkdir(parents=True, exist_ok=True)
-
-            # 2. Move any generation parameters into generation_config.json
-            base_cfg = AutoConfig.from_pretrained(input_model_id)
-            gen_cfg = GenerationConfig.from_model_config(base_cfg)
-            for k in gen_cfg.to_dict():
-                if hasattr(base_cfg, k):
-                    setattr(base_cfg, k, None)
-            base_cfg.save_pretrained(model_dir)
-            gen_cfg.save_pretrained(model_dir)
-
-            # 3. Set verbose logging via env var (no --debug flag)
-            env = os.environ.copy()
-            env["TRANSFORMERS_VERBOSITY"] = "debug"
-
-            # 4. Build and run the conversion command
             cmd = [
                 sys.executable,
                 "-m", "scripts.convert",
@@ -128,43 +108,41 @@ class ModelConverter:
                 cwd=self.config.repo_path,
                 capture_output=True,
                 text=True,
-                env=env,
+                env={},
             )
 
-            # 5. Filter out spurious warnings from stderr
-            filtered = []
-            for ln in result.stderr.splitlines():
-                if ln.startswith("Moving the following attributes"):
-                    continue
-                if "TracerWarning" in ln:
-                    continue
-                filtered.append(ln)
-            stderr = "\n".join(filtered)
-
             if result.returncode != 0:
-                return False, stderr
-            return True, stderr
+                return False, result.stderr
+
+            return True, result.stderr
 
         except Exception as e:
             return False, str(e)
 
     def upload_model(self, input_model_id: str, output_model_id: str) -> Optional[str]:
-        """Upload the converted model to Hugging Face Hub."""
-        model_folder = self.config.repo_path / "models" / input_model_id
+        """Upload the converted model to Hugging Face."""
+        model_folder_path = self.config.repo_path / "models" / input_model_id
+
         try:
             self.api.create_repo(output_model_id, exist_ok=True, private=False)
-            readme_path = model_folder / "README.md"
-            if not readme_path.exists():
-                readme_path.write_text(self.generate_readme(input_model_id))
-            self.api.upload_folder(folder_path=str(model_folder), repo_id=output_model_id)
+
+            readme_path = f"{model_folder_path}/README.md"
+
+            if not os.path.exists(readme_path):
+                with open(readme_path, "w") as file:
+                    file.write(self.generate_readme(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
-            shutil.rmtree(model_folder, ignore_errors=True)
+            shutil.rmtree(model_folder_path, ignore_errors=True)
 
-    def generate_readme(self, imi: str) -> str:
+    def generate_readme(self, imi: str):
         return (
             "---\n"
             "library_name: transformers.js\n"
@@ -173,13 +151,14 @@ class ModelConverter:
             "---\n\n"
             f"# {imi.split('/')[-1]} (ONNX)\n\n"
             f"This is an ONNX version of [{imi}](https://huggingface.co/{imi}). "
-            "Converted with attention maps and verbose export logs.\n"
+            "It was automatically converted and uploaded using "
+            "[this space](https://huggingface.co/spaces/onnx-community/convert-to-onnx).\n"
         )
 
 
 def main():
-    """Streamlit application entry point."""
-    st.write("## Convert a Hugging Face model to ONNX (with attentions & debug logs)")
+    """Main application entry point."""
+    st.write("## Convert a Hugging Face model to ONNX (with attentions)")
 
     try:
         config = Config.from_env()
@@ -187,19 +166,21 @@ def main():
         converter.setup_repository()
 
         input_model_id = st.text_input(
-            "Enter the Hugging Face model ID to convert, e.g. `EleutherAI/pythia-14m`"
+            "Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"
         )
         if not input_model_id:
             return
 
         st.text_input(
-            "Optional: Your Hugging Face write token (for uploading to your namespace).",
+            "Optional: Your Hugging Face write token. Fill it if you want to upload under your account.",
             type="password",
             key="user_hf_token",
         )
 
         if config.hf_username == input_model_id.split("/")[0]:
-            same_repo = st.checkbox("Upload ONNX weights to the same repository?")
+            same_repo = st.checkbox(
+                "Upload ONNX weights to the same repository?"
+            )
         else:
             same_repo = False
 
@@ -208,14 +189,20 @@ def main():
         if not same_repo:
             output_model_id += "-ONNX"
 
-        output_url = f"{config.hf_base_url}/{output_model_id}"
+        output_model_url = f"{config.hf_base_url}/{output_model_id}"
+
+        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
+
         st.write("Destination repository:")
-        st.code(output_url, language="plaintext")
+        st.code(output_model_url, language="plaintext")
 
-        if not st.button("Proceed", type="primary"):
+        if not st.button(label="Proceed", type="primary"):
             return
 
-        with st.spinner("Converting model…"):
+        with st.spinner("Converting model (including attention maps)…"):
             success, stderr = converter.convert_model(input_model_id)
             if not success:
                 st.error(f"Conversion failed: {stderr}")
@@ -229,12 +216,14 @@ def main():
                 st.error(f"Upload failed: {error}")
                 return
             st.success("Upload successful!")
-            st.link_button(f"Go to {output_model_id}", output_url, type="primary")
+            st.write("You can now 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: {e}")
+        st.error(f"An error occurred: {str(e)}")
 
 
 if __name__ == "__main__":
-    main()
+    main()
+