Update app.py

app.py

@@ -13,7 +13,7 @@ from huggingface_hub import HfApi, whoami
 from torch.jit import TracerWarning
 from transformers import AutoConfig, GenerationConfig
 
-# Suppress local TorchScript TracerWarnings
+# Suppress local TorchScript tracer warnings
 warnings.filterwarnings("ignore", category=TracerWarning)
 
 logging.basicConfig(level=logging.INFO)
@@ -22,6 +22,7 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class Config:
+    """Application configuration."""
     hf_token: str
     hf_username: str
     transformers_version: str = "3.5.0"
@@ -33,6 +34,7 @@ class Config:
 
     @classmethod
     def from_env(cls) -> "Config":
+        """Create config from environment variables and secrets."""
         system_token = st.secrets.get("HF_TOKEN")
         user_token = st.session_state.get("user_hf_token")
         if user_token:
@@ -48,11 +50,14 @@ class Config:
 
 
 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"
@@ -61,6 +66,7 @@ class ModelConverter:
             return "heads"
 
     def setup_repository(self) -> None:
+        """Download and setup transformers.js repo if needed."""
         if self.config.repo_path.exists():
             return
         ref_type = self._get_ref_type()
@@ -76,30 +82,39 @@ class ModelConverter:
             archive_path.unlink(missing_ok=True)
 
     def _extract_archive(self, archive_path: Path) -> None:
+        """Extract the downloaded archive."""
         import tarfile, tempfile
         with tempfile.TemporaryDirectory() as tmp_dir:
             with tarfile.open(archive_path, "r:gz") as tar:
                 tar.extractall(tmp_dir)
-            next(Path(tmp_dir).iterdir()).rename(self.config.repo_path)
+            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.
+        """
         try:
-            # Prepare model dir
+            # 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)
-            # Relocate generation params
+
+            # 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)
+                if hasattr(base_cfg, k):
+                    setattr(base_cfg, k, None)
             base_cfg.save_pretrained(model_dir)
             gen_cfg.save_pretrained(model_dir)
-            # Set verbose logging
+
+            # 3. Set verbose logging via env var (no --debug flag)
             env = os.environ.copy()
             env["TRANSFORMERS_VERBOSITY"] = "debug"
-            # Build command with debug
-            # Build conversion command
-            # Rely on TRANSFORMERS_VERBOSITY for logging; remove unsupported debug flag
+
+            # 4. Build and run the conversion command
             cmd = [
                 sys.executable,
                 "-m", "scripts.convert",
@@ -107,7 +122,6 @@ class ModelConverter:
                 "--trust_remote_code",
                 "--model_id", input_model_id,
                 "--output_attentions",
-                "--debug"
             ]
             result = subprocess.run(
                 cmd,
@@ -116,28 +130,39 @@ class ModelConverter:
                 text=True,
                 env=env,
             )
-            # Filter warnings
-            filtered = [ln for ln in result.stderr.splitlines() if not ln.startswith("Moving the following attributes") and "TracerWarning" not in ln]
+
+            # 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
+
         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
         try:
             self.api.create_repo(output_model_id, exist_ok=True, private=False)
-            readme = model_folder / "README.md"
-            if not readme.exists():
-                readme.write_text(self.generate_readme(input_model_id))
+            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)
             return None
         except Exception as e:
             return str(e)
         finally:
-            import shutil; shutil.rmtree(model_folder, ignore_errors=True)
+            import shutil
+            shutil.rmtree(model_folder, ignore_errors=True)
 
     def generate_readme(self, imi: str) -> str:
         return (
@@ -148,31 +173,68 @@ 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 debug logs and attention maps.\n"
+            "Converted with attention maps and verbose export logs.\n"
         )
 
+
 def main():
-    st.write("## Convert a Hugging Face model to ONNX (with debug)")
+    """Streamlit application entry point."""
+    st.write("## Convert a Hugging Face model to ONNX (with attentions & debug logs)")
+
     try:
         config = Config.from_env()
-        conv = ModelConverter(config)
-        conv.setup_repository()
-        input_id = st.text_input("Model ID e.g. EleutherAI/pythia-14m")
-        if not input_id: return
-        st.text_input("HF write token (optional)", type="password", key="user_hf_token")
-        same = st.checkbox("Upload to same repo?", value=False) if config.hf_username == input_id.split("/")[0] else False
-        name = input_id.split("/")[-1]; out = f"{config.hf_username}/{name}" + ("" if same else "-ONNX")
-        url = f"{config.hf_base_url}/{out}"; st.code(url)
-        if not st.button("Proceed"): return
-        with st.spinner("Converting (debug)..."):
-            ok, err = conv.convert_model(input_id)
-            if not ok: st.error(f"Conversion failed: {err}"); return
-            st.success("Conversion successful!"); st.code(err)
-        with st.spinner("Uploading..."):
-            err2 = conv.upload_model(input_id, out)
-            if err2: st.error(f"Upload failed: {err2}"); return
-            st.success("Upload successful!"); st.link_button(f"Go to {out}", url)
+        converter = ModelConverter(config)
+        converter.setup_repository()
+
+        input_model_id = st.text_input(
+            "Enter the Hugging Face model ID to convert, e.g. `EleutherAI/pythia-14m`"
+        )
+        if not input_model_id:
+            return
+
+        st.text_input(
+            "Optional: Your Hugging Face write token (for uploading to your namespace).",
+            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?")
+        else:
+            same_repo = False
+
+        model_name = input_model_id.split("/")[-1]
+        output_model_id = f"{config.hf_username}/{model_name}"
+        if not same_repo:
+            output_model_id += "-ONNX"
+
+        output_url = f"{config.hf_base_url}/{output_model_id}"
+        st.write("Destination repository:")
+        st.code(output_url, language="plaintext")
+
+        if not st.button("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
+            st.success("Upload successful!")
+            st.link_button(f"Go to {output_model_id}", output_url, type="primary")
+
     except Exception as e:
-        logger.exception(e); st.error(f"Error: {e}")
+        logger.exception("Application error")
+        st.error(f"An error occurred: {e}")
+
 
-if __name__ == "__main__": main()
+if __name__ == "__main__":
+    main()