Add `output_attentions` flag the conversion of Whisper model

app.py

@@ -7,7 +7,7 @@ import tarfile
 import shutil
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import List, Optional, Tuple
 from urllib.request import urlopen, urlretrieve
 
 import streamlit as st
@@ -128,21 +128,27 @@ class ModelConverter:
         )
 
     def convert_model(
-        self, input_model_id: str, trust_remote_code=False
+        self,
+        input_model_id: str,
+        trust_remote_code=False,
+        output_attentions=False,
     ) -> Tuple[bool, Optional[str]]:
         """Convert the model to ONNX format."""
         try:
+            extra_args: List[str] = []
             if trust_remote_code:
                 if not self.config.is_using_user_token:
                     raise Exception(
                         "Trust Remote Code requires your own HuggingFace token."
                     )
+                extra_args.append("--trust_remote_code")
 
-                result = self._run_conversion_subprocess(
-                    input_model_id, extra_args=["--trust_remote_code"]
-                )
-            else:
-                result = self._run_conversion_subprocess(input_model_id)
+            if output_attentions:
+                extra_args.append("--output_attentions")
+
+            result = self._run_conversion_subprocess(
+                input_model_id, extra_args=extra_args or None
+            )
 
             if result.returncode != 0:
                 return False, result.stderr
@@ -215,6 +221,12 @@ def main():
                 "This option should only be enabled for repositories you trust and in which you have read the code, as it will execute arbitrary code present in the model repository. When this option is enabled, you must use your own Hugging Face write token."
             )
 
+        output_attentions = False
+        if "whisper" in input_model_id.lower():
+            output_attentions = st.toggle(
+                "Whether to output attentions from the Whisper model. This is required for word-level (token) timestamps."
+            )
+
         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?"
@@ -244,7 +256,9 @@ def main():
 
         with st.spinner("Converting model..."):
             success, stderr = converter.convert_model(
-                input_model_id, trust_remote_code=trust_remote_code
+                input_model_id,
+                trust_remote_code=trust_remote_code,
+                output_attentions=output_attentions,
             )
             if not success:
                 st.error(f"Conversion failed: {stderr}")