Update app.py

app.py

@@ -93,18 +93,18 @@ class ModelConverter:
             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."""
+        """Convert the model to ONNX format, always exporting attention maps."""
         try:
+            cmd = [
+                sys.executable,
+                "-m", "scripts.convert",
+                "--quantize",
+                "--trust_remote_code",
+                "--model_id", input_model_id,
+                "--output_attentions",
+            ]
             result = subprocess.run(
-                [
-                    sys.executable,
-                    "-m",
-                    "scripts.convert",
-                    "--quantize",
-                    "--trust_remote_code",
-                    "--model_id",
-                    input_model_id,
-                ],
+                cmd,
                 cwd=self.config.repo_path,
                 capture_output=True,
                 text=True,
@@ -140,7 +140,6 @@ class ModelConverter:
             return str(e)
         finally:
             import shutil
-
             shutil.rmtree(model_folder_path, ignore_errors=True)
 
     def generate_readme(self, imi: str):
@@ -159,7 +158,7 @@ class ModelConverter:
 
 def main():
     """Main application entry point."""
-    st.write("## Convert a Hugging Face model to ONNX")
+    st.write("## Convert a Hugging Face model to ONNX (with attentions)")
 
     try:
         config = Config.from_env()
@@ -169,27 +168,24 @@ def main():
         input_model_id = st.text_input(
             "Enter the Hugging Face model ID to convert. Example: `EleutherAI/pythia-14m`"
         )
-
         if not input_model_id:
             return
 
         st.text_input(
-            f"Optional: Your Hugging Face write token. Fill it if you want to upload the model under your account.",
+            "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(
-                "Do you want to upload the ONNX weights to the same repository?"
+                "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"
 
@@ -200,29 +196,27 @@ def main():
             st.link_button(f"Go to {output_model_id}", output_model_url, type="primary")
             return
 
-        st.write(f"URL where the model will be converted and uploaded to:")
+        st.write("Destination repository:")
         st.code(output_model_url, language="plaintext")
 
         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}")
                 return
-
             st.success("Conversion successful!")
             st.code(stderr)
 
-        with st.spinner("Uploading model..."):
+        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.write("You can now go and view the model on Hugging Face!")
+            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:
@@ -232,3 +226,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+