From ba13d64bb3e5a89976a0bb3641bf2a481db55232 Mon Sep 17 00:00:00 2001
From: teleprint-me <77757836+teleprint-me@users.noreply.github.com>
Date: Fri, 17 May 2024 20:26:21 -0400
Subject: [PATCH] feat: Add utils for logging and writing when interacting with
 HuggingFaceHub

---
 gguf-py/gguf/huggingface_hub.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/gguf-py/gguf/huggingface_hub.py b/gguf-py/gguf/huggingface_hub.py
index 2e7619efb..4667ec829 100644
--- a/gguf-py/gguf/huggingface_hub.py
+++ b/gguf-py/gguf/huggingface_hub.py
@@ -1,17 +1,24 @@
 import logging
+import os
 import pathlib
 
 import requests
 
 
 class HuggingFaceHub:
-    def __init__(self, auth_token: None | str):
+    def __init__(self, auth_token: None | str, logger: None | logging.Logger):
         # Set headers if authentication is available
         if auth_token is None:
             self._headers = {}
         else:
             self._headers = {"Authorization": f"Bearer {auth_token}"}
 
+        # Set the logger
+        if logger is None:
+            logging.basicConfig(level=logging.DEBUG)
+            logger = logging.getLogger("huggingface-hub")
+        self.logger = logger
+
         # Persist across requests
         self._session = requests.Session()
 
@@ -34,6 +41,12 @@ class HuggingFaceHub:
     def base_url(self) -> str:
         return self._base_url
 
+    @staticmethod
+    def write_file(self, content: bytes, save_path: pathlib.Path) -> None:
+        with open(save_path, 'wb') as f:
+            f.write(content)
+        self.logger.info(f"File {save_path} downloaded successfully")
+
     def resolve_path(self, repo: str, file: str) -> str:
         return f"{self._base_url}/{repo}/resolve/main/{file}"