Improve-token-checking-with-whoami-api

text_classification.py

@@ -10,6 +10,7 @@ import os
 
 logger = logging.getLogger(__name__)
 HF_WRITE_TOKEN = "HF_WRITE_TOKEN"
+AUTH_CHECK_URL = "https://huggingface.co/api/whoami-v2"
 
 logger = logging.getLogger(__file__)
 
@@ -387,10 +388,9 @@ def check_hf_token_validity(hf_token):
         return False
     if not isinstance(hf_token, str):
         return False
-    # use inference api to check the token
-    payload = {"inputs": "This is a test", "options": {"use_cache": True}}
-    # use a tiny model to test the token
-    response = hf_inference_api("cross-encoder/ms-marco-TinyBERT-L-2-v2", hf_token, payload)
-    if "error" in response:
+    # use huggingface api to check the token
+    headers = {"Authorization": f"Bearer {hf_token}"}
+    response = requests.get(AUTH_CHECK_URL, headers=headers)
+    if response.status_code != 200:
         return False
     return True