codekingpro/portable-devtools

1# Copyright 2023 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""Contains a helper to get the token from machine (env variable, secret or config file)."""15 16import configparser17import logging18import os19import warnings20from pathlib import Path21from threading import Lock22 23from .. import constants24from ._runtime import is_colab_enterprise, is_google_colab25 26 27_IS_GOOGLE_COLAB_CHECKED = False28_GOOGLE_COLAB_SECRET_LOCK = Lock()29_GOOGLE_COLAB_SECRET: str | None = None30 31logger = logging.getLogger(__name__)32 33 34def get_token() -> str | None:35    """36    Get token if user is logged in.37 38    Note: in most cases, you should use [`huggingface_hub.utils.build_hf_headers`] instead. This method is only useful39          if you want to retrieve the token for other purposes than sending an HTTP request.40 41    Token is retrieved in priority from the `HF_TOKEN` environment variable. Otherwise, we read the token file located42    in the Hugging Face home folder. Returns None if user is not logged in. To log in, use [`login`] or43    `hf auth login`.44 45    Returns:46        `str` or `None`: The token, `None` if it doesn't exist.47    """48    return _get_token_from_google_colab() or _get_token_from_environment() or _get_token_from_file()49 50 51def _get_token_from_google_colab() -> str | None:52    """Get token from Google Colab secrets vault using `google.colab.userdata.get(...)`.53 54    Token is read from the vault only once per session and then stored in a global variable to avoid re-requesting55    access to the vault.56    """57    # If it's not a Google Colab or it's Colab Enterprise, fallback to environment variable or token file authentication58    if not is_google_colab() or is_colab_enterprise():59        return None60 61    # `google.colab.userdata` is not thread-safe62    # This can lead to a deadlock if multiple threads try to access it at the same time63    # (typically when using `snapshot_download`)64    # => use a lock65    # See https://github.com/huggingface/huggingface_hub/issues/1952 for more details.66    with _GOOGLE_COLAB_SECRET_LOCK:67        global _GOOGLE_COLAB_SECRET68        global _IS_GOOGLE_COLAB_CHECKED69 70        if _IS_GOOGLE_COLAB_CHECKED:  # request access only once71            return _GOOGLE_COLAB_SECRET72 73        try:74            from google.colab import userdata  # type: ignore75            from google.colab.errors import Error as ColabError  # type: ignore76        except ImportError:77            return None78 79        try:80            token = userdata.get("HF_TOKEN")81            _GOOGLE_COLAB_SECRET = _clean_token(token)82        except userdata.NotebookAccessError:83            # Means the user has a secret call `HF_TOKEN` and got a popup "please grand access to HF_TOKEN" and refused it84            # => warn user but ignore error => do not re-request access to user85            warnings.warn(86                "\nAccess to the secret `HF_TOKEN` has not been granted on this notebook."87                "\nYou will not be requested again."88                "\nPlease restart the session if you want to be prompted again."89            )90            _GOOGLE_COLAB_SECRET = None91        except userdata.SecretNotFoundError:92            # Means the user did not define a `HF_TOKEN` secret => warn93            warnings.warn(94                "\nThe secret `HF_TOKEN` does not exist in your Colab secrets."95                "\nTo authenticate with the Hugging Face Hub, create a token in your settings tab "96                "(https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session."97                "\nYou will be able to reuse this secret in all of your notebooks."98                "\nPlease note that authentication is recommended but still optional to access public models or datasets."99            )100            _GOOGLE_COLAB_SECRET = None101        except ColabError as e:102            # Something happen but we don't know what => recommend to open a GitHub issue103            warnings.warn(104                f"\nError while fetching `HF_TOKEN` secret value from your vault: '{str(e)}'."105                "\nYou are not authenticated with the Hugging Face Hub in this notebook."106                "\nIf the error persists, please let us know by opening an issue on GitHub "107                "(https://github.com/huggingface/huggingface_hub/issues/new)."108            )109            _GOOGLE_COLAB_SECRET = None110 111        _IS_GOOGLE_COLAB_CHECKED = True112        return _GOOGLE_COLAB_SECRET113 114 115def _get_token_from_environment() -> str | None:116    # `HF_TOKEN` has priority (keep `HUGGING_FACE_HUB_TOKEN` for backward compatibility)117    return _clean_token(os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN"))118 119 120def _get_token_from_file() -> str | None:121    try:122        return _clean_token(Path(constants.HF_TOKEN_PATH).read_text())123    except FileNotFoundError:124        return None125 126 127def get_stored_tokens() -> dict[str, str]:128    """129    Returns the parsed INI file containing the access tokens.130    The file is located at `HF_STORED_TOKENS_PATH`, defaulting to `~/.cache/huggingface/stored_tokens`.131    If the file does not exist, an empty dictionary is returned.132 133    Returns: `dict[str, str]`134        Key is the token name and value is the token.135    """136    tokens_path = Path(constants.HF_STORED_TOKENS_PATH)137    if not tokens_path.exists():138        stored_tokens = {}139    config = configparser.ConfigParser()140    try:141        config.read(tokens_path)142        stored_tokens = {token_name: config.get(token_name, "hf_token") for token_name in config.sections()}143    except configparser.Error as e:144        logger.error(f"Error parsing stored tokens file: {e}")145        stored_tokens = {}146    return stored_tokens147 148 149def _save_stored_tokens(stored_tokens: dict[str, str]) -> None:150    """151    Saves the given configuration to the stored tokens file.152 153    Args:154        stored_tokens (`dict[str, str]`):155            The stored tokens to save. Key is the token name and value is the token.156    """157    stored_tokens_path = Path(constants.HF_STORED_TOKENS_PATH)158 159    # Write the stored tokens into an INI file160    config = configparser.ConfigParser()161    for token_name in sorted(stored_tokens.keys()):162        config.add_section(token_name)163        config.set(token_name, "hf_token", stored_tokens[token_name])164 165    stored_tokens_path.parent.mkdir(parents=True, exist_ok=True)166    with stored_tokens_path.open("w") as config_file:167        config.write(config_file)168 169 170def _get_token_by_name(token_name: str) -> str | None:171    """172    Get the token by name.173 174    Args:175        token_name (`str`):176            The name of the token to get.177 178    Returns:179        `str` or `None`: The token, `None` if it doesn't exist.180 181    """182    stored_tokens = get_stored_tokens()183    if token_name not in stored_tokens:184        return None185    return _clean_token(stored_tokens[token_name])186 187 188def _save_token(token: str, token_name: str) -> None:189    """190    Save the given token.191 192    If the stored tokens file does not exist, it will be created.193    Args:194        token (`str`):195            The token to save.196        token_name (`str`):197            The name of the token.198    """199    tokens_path = Path(constants.HF_STORED_TOKENS_PATH)200    stored_tokens = get_stored_tokens()201    stored_tokens[token_name] = token202    _save_stored_tokens(stored_tokens)203    logger.info(f"The token `{token_name}` has been saved to {tokens_path}")204 205 206def _clean_token(token: str | None) -> str | None:207    """Clean token by removing trailing and leading spaces and newlines.208 209    If token is an empty string, return None.210    """211    if token is None:212        return None213    return token.replace("\r", "").replace("\n", "").strip() or None214