From e1b2f3864f2bbc546a8275f118edf30f8fda8ff9 Mon Sep 17 00:00:00 2001
From: Sam Havens <47401552+sam-writer@users.noreply.github.com>
Date: Wed, 8 Dec 2021 08:10:38 -0800
Subject: [PATCH] Add auth token (#19)

* add auth token

add CLI argument and a little parsing. This allows using private models from the huggingface hub.

* better parsing

* linter fix

Co-authored-by: pommedeterresautee <pommedeterresautee@msn.com>
---
 src/transformer_deploy/convert.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/transformer_deploy/convert.py b/src/transformer_deploy/convert.py
index cf4ef983..f2e950bf 100644
--- a/src/transformer_deploy/convert.py
+++ b/src/transformer_deploy/convert.py
@@ -48,6 +48,14 @@ def main():
         description="optimize and deploy transformers", formatter_class=argparse.ArgumentDefaultsHelpFormatter
     )
     parser.add_argument("-m", "--model", required=True, help="path to model or URL to Hugging Face Hub")
+    parser.add_argument(
+        "--auth-token",
+        default=None,
+        help=(
+            "HuggingFace Hub auth token. Set to `None` (default) for public models. "
+            "For private models, use `True` to use local cached token, or a string of your HF API token"
+        ),
+    )
     parser.add_argument(
         "-b",
         "--batch-size",
@@ -86,17 +94,26 @@ def main():
 
     torch.manual_seed(args.seed)
 
+    if isinstance(args.auth_token, str) and args.auth_token.lower() in ["true", "t"]:
+        auth_token = True
+    elif isinstance(args.auth_token, str):
+        auth_token = args.auth_token
+    else:
+        auth_token = None
+
     Path(args.output).mkdir(parents=True, exist_ok=True)
     onnx_model_path = os.path.join(args.output, "model-original.onnx")
     onnx_optim_fp16_path = os.path.join(args.output, "model.onnx")
     tensorrt_path = os.path.join(args.output, "model.plan")
 
     assert torch.cuda.is_available(), "CUDA is not available. Please check your CUDA installation"
-    tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(args.model)
+    tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(args.model, use_auth_token=auth_token)
     input_names: List[str] = tokenizer.model_input_names
     logging.info(f"axis: {input_names}")
     include_token_ids = "token_type_ids" in input_names
-    model_pytorch: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(args.model)
+    model_pytorch: PreTrainedModel = AutoModelForSequenceClassification.from_pretrained(
+        args.model, use_auth_token=auth_token
+    )
     model_pytorch.cuda()
     model_pytorch.eval()