Update app.py

app.py

@@ -2,6 +2,7 @@ import logging
 import os
 import subprocess
 import sys
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, Tuple
@@ -9,6 +10,11 @@ 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 TorchScript tracer warnings globally
+warnings.filterwarnings("ignore", category=TracerWarning)
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -88,13 +94,31 @@ class ModelConverter:
         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, always exporting attention maps."""
+        """
+        Convert the model to ONNX format, always exporting attention maps.
+        Also relocate any generation parameters into generation_config.json
+        and suppress JIT tracer warnings.
+        """
         try:
+            # 1. Clone or prepare a local copy of the model to adjust configs
+            model_dir = self.config.repo_path / "models" / input_model_id
+            model_dir.mkdir(parents=True, exist_ok=True)
+
+            # 2. Load and relocate generation parameters
+            config = AutoConfig.from_pretrained(input_model_id)
+            gen_cfg = GenerationConfig.from_model_config(config)
+            # Remove generation-specific keys from model config
+            for key in gen_cfg.to_dict().keys():
+                if hasattr(config, key):
+                    setattr(config, key, None)
+            config.save_pretrained(model_dir)
+            gen_cfg.save_pretrained(model_dir)
+
+            # 3. Build the conversion command
             cmd = [
                 sys.executable,
                 "-m", "scripts.convert",
@@ -103,12 +127,14 @@ class ModelConverter:
                 "--model_id", input_model_id,
                 "--output_attentions",
             ]
+
+            # 4. Run the conversion
             result = subprocess.run(
                 cmd,
                 cwd=self.config.repo_path,
                 capture_output=True,
                 text=True,
-                env={},
+                env=os.environ.copy(),
             )
 
             if result.returncode != 0:
@@ -127,13 +153,13 @@ class ModelConverter:
             self.api.create_repo(output_model_id, exist_ok=True, private=False)
 
             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
+                folder_path=str(model_folder_path),
+                repo_id=output_model_id
             )
             return None
         except Exception as e:
@@ -142,7 +168,7 @@ class ModelConverter:
             import shutil
             shutil.rmtree(model_folder_path, ignore_errors=True)
 
-    def generate_readme(self, imi: str):
+    def generate_readme(self, imi: str) -> str:
         return (
             "---\n"
             "library_name: transformers.js\n"
@@ -178,9 +204,7 @@ def main():
         )
 
         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
 
@@ -226,4 +250,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-