ZeroCommand commited on
Commit
db8ac73
1 Parent(s): 44e6352

Improve-token-checking-with-whoami-api (#110)

Browse files

- use whoami api to verify token (a80c377d57bf235bbffa0742914f46556af45d02)

Files changed (1) hide show
  1. text_classification.py +5 -5
text_classification.py CHANGED
@@ -10,6 +10,7 @@ import os
10
 
11
  logger = logging.getLogger(__name__)
12
  HF_WRITE_TOKEN = "HF_WRITE_TOKEN"
 
13
 
14
  logger = logging.getLogger(__file__)
15
 
@@ -387,10 +388,9 @@ def check_hf_token_validity(hf_token):
387
  return False
388
  if not isinstance(hf_token, str):
389
  return False
390
- # use inference api to check the token
391
- payload = {"inputs": "This is a test", "options": {"use_cache": True}}
392
- # use a tiny model to test the token
393
- response = hf_inference_api("cross-encoder/ms-marco-TinyBERT-L-2-v2", hf_token, payload)
394
- if "error" in response:
395
  return False
396
  return True
 
10
 
11
  logger = logging.getLogger(__name__)
12
  HF_WRITE_TOKEN = "HF_WRITE_TOKEN"
13
+ AUTH_CHECK_URL = "https://huggingface.co/api/whoami-v2"
14
 
15
  logger = logging.getLogger(__file__)
16
 
 
388
  return False
389
  if not isinstance(hf_token, str):
390
  return False
391
+ # use huggingface api to check the token
392
+ headers = {"Authorization": f"Bearer {hf_token}"}
393
+ response = requests.get(AUTH_CHECK_URL, headers=headers)
394
+ if response.status_code != 200:
 
395
  return False
396
  return True