feat: Add prototype for requesting vocab related files

This commit is contained in:
teleprint-me 2024-05-17 21:08:39 -04:00
parent 98cf788990
commit 4790f76740
No known key found for this signature in database
GPG key ID: B0D11345E65C4D48

View file

@ -57,3 +57,46 @@ class HuggingFaceHub:
self.logger.info(f"Response status was {response.status_code}") self.logger.info(f"Response status was {response.status_code}")
response.raise_for_status() response.raise_for_status()
return response return response
class HFTokenizerRequest:
def __init__(
self,
dl_path: pathlib.Path,
auth_token: str,
logger: None | logging.Logger
):
self._hub = HuggingFaceHub(auth_token, logger)
self._models = MODEL_REPOS
if dl_path is None:
self._download_path = pathlib.Path("models/tokenizers")
else:
self._download_path = dl_path
self._files = ["config.json", "tokenizer_config.json", "tokenizer.json"]
@property
def hub(self) -> HuggingFaceHub:
return self._hub
@property
def models(self) -> list[dict[str, object]]:
return self._models
@property
def download_path(self) -> pathlib.Path:
return self._download_path
@download_path.setter
def download_path(self, value: pathlib.Path):
self._download_path = value
@property
def files(self) -> list[str]:
return self._files
def download_file_with_auth(self, repo, file, directory):
response = self.hub.download_file(repo, file)
os.makedirs(os.path.dirname(directory), exist_ok=True)
self.hub.write_file(response.content, directory)