Propagate token through to all ST/HF_hub functions

README.md

@@ -10,6 +10,10 @@ pinned: false
 license: apache-2.0
 short_description: Export Sentence Transformer models to accelerated backends
 hf_oauth: true
+hf_oauth_scopes:
+- manage-repos
+- write-repos
+- write-discussions
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,15 +1,17 @@
 from enum import Enum
+from functools import partial
 from pathlib import Path
-from typing import Tuple
+from typing import Optional, Tuple
 import gradio as gr
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
+import huggingface_hub
 from sentence_transformers import SentenceTransformer
 from sentence_transformers import (
     export_dynamic_quantized_onnx_model as st_export_dynamic_quantized_onnx_model,
     export_optimized_onnx_model as st_export_optimized_onnx_model,
     export_static_quantized_openvino_model as st_export_static_quantized_openvino_model,
 )
-from huggingface_hub import model_info, upload_folder, whoami, get_repo_discussions, list_repo_commits, HfFileSystem
+from huggingface_hub import model_info, upload_folder, get_repo_discussions, list_repo_commits, HfFileSystem
 from huggingface_hub.errors import RepositoryNotFoundError
 from optimum.intel import OVQuantizationConfig
 from tempfile import TemporaryDirectory
@@ -78,7 +80,7 @@ def export_to_torch(model_id, create_pr, output_model_id):
     )
 
 
-def export_to_onnx(model_id: str, create_pr: bool, output_model_id: str):
+def export_to_onnx(model_id: str, create_pr: bool, output_model_id: str, token: Optional[str] = None) -> None:
     if does_file_glob_exist(output_model_id, "**/model.onnx"):
         raise FileExistsError("An ONNX model already exists in the repository")
 
@@ -91,6 +93,7 @@ def export_to_onnx(model_id: str, create_pr: bool, output_model_id: str):
             repo_id=output_model_id,
             commit_message=commit_message,
             create_pr=create_pr,
+            token=token,
         )
     else:
         with TemporaryDirectory() as tmp_dir:
@@ -133,6 +136,7 @@ print(similarities)
                 commit_message=commit_message,
                 commit_description=commit_description if create_pr else None,
                 create_pr=create_pr,
+                token=token,
             )
 
 def export_to_onnx_snippet(model_id: str, create_pr: bool, output_model_id: str) -> str:
@@ -175,7 +179,7 @@ similarities = model.similarity(embeddings, embeddings)
 
 
 def export_to_onnx_dynamic_quantization(
-    model_id: str, create_pr: bool, output_model_id: str, onnx_quantization_config: str
+    model_id: str, create_pr: bool, output_model_id: str, onnx_quantization_config: str, token: Optional[str] = None
 ) -> None:
     if does_file_glob_exist(output_model_id, f"onnx/model_qint8_{onnx_quantization_config}.onnx"):
         raise FileExistsError("The quantized ONNX model already exists in the repository")
@@ -183,8 +187,11 @@ def export_to_onnx_dynamic_quantization(
     model = SentenceTransformer(model_id, backend="onnx")
 
     if not create_pr and is_new_model(output_model_id):
-        model.push_to_hub(repo_id=output_model_id)
+        model.push_to_hub(repo_id=output_model_id, token=token)
 
+    # Monkey-patch the upload_folder function to include the token, as it's not used in export_dynamic_quantized_onnx_model
+    original_upload_folder = huggingface_hub.upload_folder
+    huggingface_hub.upload_folder = partial(original_upload_folder, token=token)
     try:
         st_export_dynamic_quantized_onnx_model(
             model,
@@ -203,6 +210,8 @@ def export_to_onnx_dynamic_quantization(
             push_to_hub=True,
             create_pr=create_pr,
         )
+    finally:
+        huggingface_hub.upload_folder = original_upload_folder
 
 def export_to_onnx_dynamic_quantization_snippet(
     model_id: str, create_pr: bool, output_model_id: str, onnx_quantization_config: str
@@ -249,22 +258,28 @@ embeddings = model.encode(["The weather is lovely today.", "It's so sunny outsid
 similarities = model.similarity(embeddings, embeddings)
 """
 
-def export_to_onnx_optimization(model_id: str, create_pr: bool, output_model_id: str, onnx_optimization_config: str) -> None:
+def export_to_onnx_optimization(model_id: str, create_pr: bool, output_model_id: str, onnx_optimization_config: str, token: Optional[str] = None) -> None:
     if does_file_glob_exist(output_model_id, f"onnx/model_{onnx_optimization_config}.onnx"):
         raise FileExistsError("The optimized ONNX model already exists in the repository")
 
     model = SentenceTransformer(model_id, backend="onnx")
 
     if not create_pr and is_new_model(output_model_id):
-        model.push_to_hub(repo_id=output_model_id)
+        model.push_to_hub(repo_id=output_model_id, token=token)
 
-    st_export_optimized_onnx_model(
-        model,
-        optimization_config=onnx_optimization_config,
-        model_name_or_path=output_model_id,
-        push_to_hub=True,
-        create_pr=create_pr,
-    )
+    # Monkey-patch the upload_folder function to include the token, as it's not used in export_optimized_onnx_model
+    original_upload_folder = huggingface_hub.upload_folder
+    huggingface_hub.upload_folder = partial(original_upload_folder, token=token)
+    try:
+        st_export_optimized_onnx_model(
+            model,
+            optimization_config=onnx_optimization_config,
+            model_name_or_path=output_model_id,
+            push_to_hub=True,
+            create_pr=create_pr,
+        )
+    finally:
+        huggingface_hub.upload_folder = original_upload_folder
 
 def export_to_onnx_optimization_snippet(model_id: str, create_pr: bool, output_model_id: str, onnx_optimization_config: str) -> str:
     return """\
@@ -310,7 +325,7 @@ similarities = model.similarity(embeddings, embeddings)
 """
 
 
-def export_to_openvino(model_id: str, create_pr: bool, output_model_id: str) -> None:
+def export_to_openvino(model_id: str, create_pr: bool, output_model_id: str, token: Optional[str] = None) -> None:
     if does_file_glob_exist(output_model_id, "**/openvino_model.xml"):
         raise FileExistsError("The OpenVINO model already exists in the repository")
 
@@ -323,6 +338,7 @@ def export_to_openvino(model_id: str, create_pr: bool, output_model_id: str) ->
             repo_id=output_model_id,
             commit_message=commit_message,
             create_pr=create_pr,
+            token=token,
         )
     else:
         with TemporaryDirectory() as tmp_dir:
@@ -365,6 +381,7 @@ print(similarities)
                 commit_message=commit_message,
                 commit_description=commit_description if create_pr else None,
                 create_pr=create_pr,
+                token=token,
             )
 
 def export_to_openvino_snippet(model_id: str, create_pr: bool, output_model_id: str) -> str:
@@ -412,6 +429,7 @@ def export_to_openvino_static_quantization(
     ov_quant_dataset_split: str,
     ov_quant_dataset_column_name: str,
     ov_quant_dataset_num_samples: int,
+    token: Optional[str] = None,
 ) -> None:
     if does_file_glob_exist(output_model_id, "openvino/openvino_model_qint8_quantized.xml"):
         raise FileExistsError("The quantized OpenVINO model already exists in the repository")
@@ -419,21 +437,27 @@ def export_to_openvino_static_quantization(
     model = SentenceTransformer(model_id, backend="openvino")
 
     if not create_pr and is_new_model(output_model_id):
-        model.push_to_hub(repo_id=output_model_id)
-
-    st_export_static_quantized_openvino_model(
-        model,
-        quantization_config=OVQuantizationConfig(
-            num_samples=ov_quant_dataset_num_samples,
-        ),
-        model_name_or_path=output_model_id,
-        dataset_name=ov_quant_dataset_name,
-        dataset_config_name=ov_quant_dataset_subset,
-        dataset_split=ov_quant_dataset_split,
-        column_name=ov_quant_dataset_column_name,
-        push_to_hub=True,
-        create_pr=create_pr,
-    )
+        model.push_to_hub(repo_id=output_model_id, token=token)
+
+    # Monkey-patch the upload_folder function to include the token, as it's not used in export_static_quantized_openvino_model
+    original_upload_folder = huggingface_hub.upload_folder
+    huggingface_hub.upload_folder = partial(original_upload_folder, token=token)
+    try:
+        st_export_static_quantized_openvino_model(
+            model,
+            quantization_config=OVQuantizationConfig(
+                num_samples=ov_quant_dataset_num_samples,
+            ),
+            model_name_or_path=output_model_id,
+            dataset_name=ov_quant_dataset_name,
+            dataset_config_name=ov_quant_dataset_subset,
+            dataset_split=ov_quant_dataset_split,
+            column_name=ov_quant_dataset_column_name,
+            push_to_hub=True,
+            create_pr=create_pr,
+        )
+    finally:
+        huggingface_hub.upload_folder = original_upload_folder
 
 def export_to_openvino_static_quantization_snippet(
     model_id: str,
@@ -505,7 +529,11 @@ def on_submit(
     ov_quant_dataset_column_name,
     ov_quant_dataset_num_samples,
     inference_snippet: str,
+    oauth_token: Optional[gr.OAuthToken] = None,
+    profile: Optional[gr.OAuthProfile] = None,
 ):
+    if oauth_token is None or profile is None:
+        return "Commit or PR url:<br>...", inference_snippet, gr.Textbox("Please sign in with Hugging Face to use this Space", visible=True)
 
     if not model_id:
         return "Commit or PR url:<br>...", inference_snippet, gr.Textbox("Please enter a model ID", visible=True)
@@ -514,26 +542,23 @@ def on_submit(
         return "Commit or PR url:<br>...", inference_snippet, gr.Textbox("The source model must have a Sentence Transformers tag", visible=True)
 
     if output_model_id and "/" not in output_model_id:
-        try:
-            output_model_id = f"{whoami()['name']}/{output_model_id}"
-        except Exception:
-            return "Commit or PR url:<br>...", inference_snippet, gr.Textbox("You might be signed in with Hugging Face to use this Space", visible=True)
+        output_model_id = f"{profile.name}/{output_model_id}"
 
     output_model_id = output_model_id if not create_pr else model_id
 
     try:
         if backend == Backend.ONNX.value:
-            export_to_onnx(model_id, create_pr, output_model_id)
+            export_to_onnx(model_id, create_pr, output_model_id, token=oauth_token.token)
         elif backend == Backend.ONNX_DYNAMIC_QUANTIZATION.value:
             export_to_onnx_dynamic_quantization(
-                model_id, create_pr, output_model_id, onnx_quantization_config
+                model_id, create_pr, output_model_id, onnx_quantization_config, token=oauth_token.token
             )
         elif backend == Backend.ONNX_OPTIMIZATION.value:
             export_to_onnx_optimization(
-                model_id, create_pr, output_model_id, onnx_optimization_config
+                model_id, create_pr, output_model_id, onnx_optimization_config, token=oauth_token.token
             )
         elif backend == Backend.OPENVINO.value:
-            export_to_openvino(model_id, create_pr, output_model_id)
+            export_to_openvino(model_id, create_pr, output_model_id, token=oauth_token.token)
         elif backend == Backend.OPENVINO_STATIC_QUANTIZATION.value:
             export_to_openvino_static_quantization(
                 model_id,
@@ -544,6 +569,7 @@ def on_submit(
                 ov_quant_dataset_split,
                 ov_quant_dataset_column_name,
                 ov_quant_dataset_num_samples,
+                token=oauth_token.token,
             )
     except FileExistsError as exc:
         return "Commit or PR url:<br>...", inference_snippet, gr.Textbox(str(exc), visible=True)
@@ -572,15 +598,17 @@ def on_change(
     ov_quant_dataset_split,
     ov_quant_dataset_column_name,
     ov_quant_dataset_num_samples,
+    oauth_token: Optional[gr.OAuthToken] = None,
+    profile: Optional[gr.OAuthProfile] = None,
 ) -> str:
+    if oauth_token is None or profile is None:
+        return "", "", "", gr.Textbox("Please sign in with Hugging Face to use this Space", visible=True)
+
     if not model_id:
         return "", "", "", gr.Textbox("Please enter a model ID", visible=True)
     
     if output_model_id and "/" not in output_model_id:
-        try:
-            output_model_id = f"{whoami()['name']}/{output_model_id}"
-        except Exception:
-            return "", "", "", gr.Textbox("You might be signed in with Hugging Face to use this Space", visible=True)
+        output_model_id = f"{profile.username}/{output_model_id}"
 
     output_model_id = output_model_id if not create_pr else model_id