toolbox/hf_netdisk/hf_netdisk_client.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse
import json
import logging
import shutil
import tempfile
from pathlib import Path

import requests
from huggingface_hub import HfApi, login

from project_settings import project_path, environment
from toolbox.design_patterns.singleton import ParamsSingleton


logger = logging.getLogger("toolbox")


class HfNetdiskClient(ParamsSingleton):

    def __init__(self, hf_token: str = None):
        if not self._initialized:
            self.hf_token = hf_token

            self.hf_api = HfApi(token=hf_token)

            self._initialized = True

    def check_login(self):
        flag = False
        try:
            user_info = self.hf_api.whoami()
            user_info = json.dumps(user_info, ensure_ascii=False)
            flag = True
        except requests.exceptions.HTTPError as error:
            user_info = f"whoami failed; error type: {type(error)}, error text: {str(error)}"
        except Exception as error:
            user_info = f"whoami failed; error type: {type(error)}, error text: {str(error)}"

        logger.info(f"login flag: {flag}, user_info: {user_info}")
        return flag


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--key_of_credentials",
        default="hf_miyuki_write_token",
        type=str,
    )
    args = parser.parse_args()
    return args


def main():
    args = get_args()

    import log
    from project_settings import environment, project_path, log_directory, time_zone_info

    log.setup_size_rotating(log_directory=log_directory, tz_info=time_zone_info)

    hf_token = environment.get(key=args.key_of_credentials)

    client = HfNetdiskClient(
        hf_token=hf_token,
    )

    client.check_login()


    return


if __name__ == "__main__":
    main()