273 lines
9.7 KiB
Python
273 lines
9.7 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
import pathlib
|
|
from hashlib import sha256
|
|
|
|
import requests
|
|
from huggingface_hub import login, model_info
|
|
from sentencepiece import SentencePieceProcessor
|
|
|
|
from .constants import (
|
|
GPT_PRE_TOKENIZER_DEFAULT,
|
|
HF_TOKENIZER_BPE_FILES,
|
|
HF_TOKENIZER_SPM_FILES,
|
|
ModelFileExtension,
|
|
NormalizerType,
|
|
PreTokenizerType,
|
|
VocabType,
|
|
)
|
|
|
|
|
|
class HFHubBase:
|
|
def __init__(
|
|
self,
|
|
model_path: None | str | pathlib.Path,
|
|
logger: None | logging.Logger,
|
|
):
|
|
# Set the model path
|
|
if model_path is None:
|
|
model_path = "models"
|
|
self._model_path = model_path
|
|
|
|
# Set the logger
|
|
if logger is None:
|
|
logger = logging.getLogger(__name__)
|
|
self.logger = logger
|
|
|
|
@property
|
|
def model_path(self) -> pathlib.Path:
|
|
return pathlib.Path(self._model_path)
|
|
|
|
@model_path.setter
|
|
def model_path(self, value: pathlib.Path):
|
|
self._model_path = value
|
|
|
|
def write_file(self, content: bytes, file_path: pathlib.Path) -> None:
|
|
with open(file_path, "wb") as file:
|
|
file.write(content)
|
|
self.logger.info(f"Wrote {len(content)} bytes to {file_path} successfully")
|
|
|
|
|
|
class HFHubRequest(HFHubBase):
|
|
def __init__(
|
|
self,
|
|
auth_token: None | str,
|
|
model_path: None | str | pathlib.Path,
|
|
logger: None | logging.Logger,
|
|
):
|
|
super().__init__(model_path, logger)
|
|
|
|
# Set headers if authentication is available
|
|
if auth_token is None:
|
|
self._headers = None
|
|
else:
|
|
self._headers = {"Authorization": f"Bearer {auth_token}"}
|
|
|
|
# Persist across requests
|
|
self._session = requests.Session()
|
|
|
|
# This is read-only
|
|
self._base_url = "https://huggingface.co"
|
|
|
|
# NOTE: Required for getting model_info
|
|
login(auth_token, add_to_git_credential=True)
|
|
|
|
@property
|
|
def headers(self) -> str:
|
|
return self._headers
|
|
|
|
@property
|
|
def session(self) -> requests.Session:
|
|
return self._session
|
|
|
|
@property
|
|
def base_url(self) -> str:
|
|
return self._base_url
|
|
|
|
@staticmethod
|
|
def list_remote_files(model_repo: str) -> list[str]:
|
|
# NOTE: Request repository metadata to extract remote filenames
|
|
return [x.rfilename for x in model_info(model_repo).siblings]
|
|
|
|
def list_filtered_remote_files(
|
|
self, model_repo: str, file_extension: ModelFileExtension
|
|
) -> list[str]:
|
|
model_files = []
|
|
self.logger.info(f"Repo:{model_repo}")
|
|
self.logger.debug(f"FileExtension:{file_extension.value}")
|
|
for filename in HFHubRequest.list_remote_files(model_repo):
|
|
suffix = pathlib.Path(filename).suffix
|
|
self.logger.debug(f"Suffix: {suffix}")
|
|
if suffix == file_extension.value:
|
|
self.logger.info(f"File: {filename}")
|
|
model_files.append(filename)
|
|
return model_files
|
|
|
|
def resolve_url(self, repo: str, filename: str) -> str:
|
|
return f"{self._base_url}/{repo}/resolve/main/{filename}"
|
|
|
|
def get_response(self, url: str) -> requests.Response:
|
|
response = self._session.get(url, headers=self.headers)
|
|
self.logger.info(f"Response status was {response.status_code}")
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
|
|
class HFHubTokenizer(HFHubBase):
|
|
def __init__(
|
|
self,
|
|
model_path: None | str | pathlib.Path,
|
|
logger: None | logging.Logger,
|
|
):
|
|
super().__init__(model_path, logger)
|
|
|
|
@staticmethod
|
|
def list_vocab_files(vocab_type: VocabType) -> tuple[str]:
|
|
if vocab_type == VocabType.SPM.value:
|
|
return HF_TOKENIZER_SPM_FILES
|
|
# NOTE: WPM and BPE are equivalent
|
|
return HF_TOKENIZER_BPE_FILES
|
|
|
|
@property
|
|
def default_pre_tokenizer(self) -> tuple[str, ...]:
|
|
return GPT_PRE_TOKENIZER_DEFAULT
|
|
|
|
def config(self, model_repo: str) -> dict[str, object]:
|
|
path = self.model_path / model_repo / "config.json"
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
def tokenizer_model(self, model_repo: str) -> SentencePieceProcessor:
|
|
path = self.model_path / model_repo / "tokenizer.model"
|
|
processor = SentencePieceProcessor()
|
|
processor.LoadFromFile(path.read_bytes())
|
|
return processor
|
|
|
|
def tokenizer_config(self, model_repo: str) -> dict[str, object]:
|
|
path = self.model_path / model_repo / "tokenizer_config.json"
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
def tokenizer_json(self, model_repo: str) -> dict[str, object]:
|
|
path = self.model_path / model_repo / "tokenizer.json"
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
def get_normalizer(self, model_repo: str) -> None | dict[str, object]:
|
|
normalizer = self.tokenizer_json(model_repo).get("normalizer", dict())
|
|
if normalizer:
|
|
self.logger.info(f"JSON:Normalizer: {json.dumps(normalizer, indent=2)}")
|
|
else:
|
|
self.logger.warn(f"WARN:Normalizer: {normalizer}")
|
|
return normalizer
|
|
|
|
def get_pre_tokenizer(self, model_repo: str) -> None | dict[str, object]:
|
|
pre_tokenizer = self.tokenizer_json(model_repo).get("pre_tokenizer")
|
|
if pre_tokenizer:
|
|
self.logger.info(
|
|
f"JSON:PreTokenizer: {json.dumps(pre_tokenizer, indent=2)}"
|
|
)
|
|
return pre_tokenizer
|
|
else:
|
|
self.logger.warn(f"WARN:PreTokenizer: {pre_tokenizer}")
|
|
return pre_tokenizer
|
|
|
|
def get_added_tokens(self, model_repo: str) -> None | list[dict[str, object]]:
|
|
added_tokens = self.tokenizer_json(model_repo).get("added_tokens", list())
|
|
if added_tokens:
|
|
self.logger.info(f"JSON:AddedTokens: {json.dumps(added_tokens, indent=2)}")
|
|
else:
|
|
self.logger.warn(f"WARN:PreTokenizer: {added_tokens}")
|
|
return added_tokens
|
|
|
|
def get_pre_tokenizer_json_hash(self, model_repo: str) -> None | str:
|
|
tokenizer = self.tokenizer_json(model_repo)
|
|
tokenizer_path = self.model_path / model_repo / "tokenizer.json"
|
|
if tokenizer.get("pre_tokenizer"):
|
|
sha256sum = sha256(str(tokenizer.get("pre_tokenizer")).encode()).hexdigest()
|
|
else:
|
|
return
|
|
self.logger.info(f"Hashed '{tokenizer_path}' as {sha256sum}")
|
|
return sha256sum
|
|
|
|
def get_tokenizer_json_hash(self, model_repo: str) -> str:
|
|
tokenizer = self.tokenizer_json(model_repo)
|
|
tokenizer_path = self.model_path / model_repo / "tokenizer.json"
|
|
sha256sum = sha256(str(tokenizer).encode()).hexdigest()
|
|
self.logger.info(f"Hashed '{tokenizer_path}' as {sha256sum}")
|
|
return sha256sum
|
|
|
|
def log_tokenizer_json_info(self, model_repo: str) -> None:
|
|
self.logger.info(f"{model_repo}")
|
|
tokenizer = self.tokenizer_json(model_repo)
|
|
for k, v in tokenizer.items():
|
|
if k not in ["added_tokens", "model"]:
|
|
self.logger.info(f"{k}:{json.dumps(v, indent=2)}")
|
|
if k == "model":
|
|
for x, y in v.items():
|
|
if x not in ["vocab", "merges"]:
|
|
self.logger.info(f"{k}:{x}:{json.dumps(y, indent=2)}")
|
|
|
|
|
|
class HFHubModel(HFHubBase):
|
|
def __init__(
|
|
self,
|
|
auth_token: None | str,
|
|
model_path: None | str | pathlib.Path,
|
|
logger: None | logging.Logger,
|
|
):
|
|
super().__init__(model_path, logger)
|
|
|
|
self._request = HFHubRequest(auth_token, model_path, logger)
|
|
self._tokenizer = HFHubTokenizer(model_path, logger)
|
|
|
|
@property
|
|
def request(self) -> HFHubRequest:
|
|
return self._request
|
|
|
|
@property
|
|
def tokenizer(self) -> HFHubTokenizer:
|
|
return self._tokenizer
|
|
|
|
def _request_single_file(
|
|
self, model_repo: str, file_name: str, file_path: pathlib.Path
|
|
) -> bool:
|
|
# NOTE: Consider optional `force` parameter if files need to be updated.
|
|
# e.g. The model creator updated the vocabulary to resolve an issue or add a feature.
|
|
if file_path.exists():
|
|
self.logger.info(f"skipped - downloaded {file_path} exists already.")
|
|
return False
|
|
|
|
# NOTE: Do not use bare exceptions! They mask issues!
|
|
# Allow the exception to occur or explicitly handle it.
|
|
try:
|
|
self.logger.info(f"Downloading '{file_name}' from {model_repo}")
|
|
resolved_url = self.request.resolve_url(model_repo, file_name)
|
|
response = self.request.get_response(resolved_url)
|
|
self.write_file(response.content, file_path)
|
|
self.logger.info(f"Model file successfully saved to {file_path}")
|
|
return True
|
|
except requests.exceptions.HTTPError as e:
|
|
self.logger.error(f"Error while downloading '{file_name}': {str(e)}")
|
|
return False
|
|
|
|
def _request_listed_files(self, model_repo: str, remote_files: list[str]) -> None:
|
|
for file_name in remote_files:
|
|
dir_path = self.model_path / model_repo
|
|
os.makedirs(dir_path, exist_ok=True)
|
|
self._request_single_file(model_repo, file_name, dir_path / file_name)
|
|
|
|
def download_model_files(
|
|
self, model_repo: str, file_extension: ModelFileExtension
|
|
) -> None:
|
|
filtered_files = self.request.list_filtered_remote_files(
|
|
model_repo, file_extension
|
|
)
|
|
self._request_listed_files(model_repo, filtered_files)
|
|
|
|
def download_all_vocab_files(self, model_repo: str, vocab_type: VocabType) -> None:
|
|
vocab_files = self.tokenizer.list_vocab_files(vocab_type)
|
|
self._request_listed_files(model_repo, vocab_files)
|
|
|
|
def download_all_model_files(self, model_repo: str) -> None:
|
|
all_files = self.request.list_remote_files(model_repo)
|
|
self._request_listed_files(model_repo, all_files)
|