[hf_hub] Token login (#3739)

* login on token

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup old code

* safer imports

* cleanup

* Return token after login

* correct return types

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply suggestion from @danielhanchen

* add back imports

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* finish return token

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

unsloth/models/_utils.py

@@ -72,6 +72,7 @@ __all__ = [
     "patch_hf_quantizer",
     "verify_fp8_support_if_applicable",
     "_get_inference_mode_context_manager",
+    "hf_login",
 ]
 
 import torch
@@ -2344,3 +2345,23 @@ def _get_inference_mode_context_manager(model: torch.nn.Module):
         return torch.no_grad()
     else:
         return torch.inference_mode()
+
+
+def hf_login(token: Optional[str] = None) -> Optional[str]:
+    if token is None:
+        try:
+            from huggingface_hub import get_token
+
+            token = get_token()
+            if token is None:
+                return None
+        except:
+            return None
+    try:
+        from huggingface_hub import login
+
+        login(token = token)
+        return token
+    except Exception as e:
+        logger.info(f"Failed to login to huggingface using token with error: {e}")
+    return token

unsloth/models/llama.py

@@ -2130,8 +2130,7 @@ class FastLlamaModel:
                     "Unsloth: `unsloth_vllm_standby` is True, but  environment variable `UNSLOTH_VLLM_STANDBY` is not set to 1!"
                 )
 
-        if token is None:
-            token = get_token()
+        token = hf_login(token)
         if model_patcher is None:
             model_patcher = FastLlamaModel
         SUPPORTS_BFLOAT16 = is_bfloat16_supported()

unsloth/models/loader.py

@@ -20,6 +20,7 @@ from ._utils import (
     HAS_FLASH_ATTENTION_SOFTCAPPING,
     USE_MODELSCOPE,
     get_transformers_model_type,
+    hf_login,
 )
 from .granite import FastGraniteModel
 from .llama import FastLlamaModel, logger
@@ -151,15 +152,7 @@ class FastLanguageModel(FastLlamaModel):
         **kwargs,
     ):
         # Login to allow private models
-        if token is None:
-            token = get_token()
-        if token is not None:
-            try:
-                from huggingface_hub import login
-
-                login(token = token)
-            except:
-                pass
+        token = hf_login(token)
         if load_in_8bit or full_finetuning or qat_scheme is not None:
             return FastModel.from_pretrained(
                 model_name = model_name,
@@ -195,8 +188,6 @@ class FastLanguageModel(FastLlamaModel):
                 **kwargs,
             )
 
-        if token is None:
-            token = get_token()
         if isinstance(dtype, str) and dtype in ["float16", "bfloat16"]:
             dtype = getattr(torch, dtype)
         assert (
@@ -682,16 +673,8 @@ class FastModel(FastBaseModel):
         *args,
         **kwargs,
     ):
-        if token is None:
-            token = get_token()
         # Login to allow private models
-        if token is not None:
-            try:
-                from huggingface_hub import login
-
-                login(token = token)
-            except:
-                pass
+        token = hf_login(token)
         if whisper_language is not None:
             assert type(whisper_language) is str
         if whisper_task is not None:

unsloth/models/vision.py

@@ -390,8 +390,7 @@ class FastBaseModel:
                 "Unsloth: WARNING `trust_remote_code` is True.\n"
                 "Are you certain you want to do remote code execution?"
             )
-        if token is None:
-            token = get_token()
+        token = hf_login(token)
         SUPPORTS_BFLOAT16 = is_bfloat16_supported()
 
         if DEVICE_TYPE == "cuda":