refactor: Simplify the huggingface hub api to enable flexible model requests
This commit is contained in:
parent
6c9ac0fc52
commit
64096942ce
1 changed files with 36 additions and 41 deletions
|
@ -5,8 +5,8 @@ import pathlib
|
||||||
from hashlib import sha256
|
from hashlib import sha256
|
||||||
|
|
||||||
import requests
|
import requests
|
||||||
|
from huggingface_hub import login, model_info
|
||||||
from sentencepiece import SentencePieceProcessor
|
from sentencepiece import SentencePieceProcessor
|
||||||
from transformers import AutoTokenizer
|
|
||||||
|
|
||||||
from .constants import (
|
from .constants import (
|
||||||
GPT_PRE_TOKENIZER_DEFAULT,
|
GPT_PRE_TOKENIZER_DEFAULT,
|
||||||
|
@ -69,7 +69,7 @@ class HFHubRequest:
|
||||||
return response
|
return response
|
||||||
|
|
||||||
|
|
||||||
class HFHubBase:
|
class HFHub:
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
model_path: None | str | pathlib.Path,
|
model_path: None | str | pathlib.Path,
|
||||||
|
@ -104,7 +104,7 @@ class HFHubBase:
|
||||||
self._model_path = value
|
self._model_path = value
|
||||||
|
|
||||||
|
|
||||||
class HFTokenizer(HFHubBase):
|
class HFTokenizer(HFHub):
|
||||||
def __init__(self, model_path: str, auth_token: str, logger: logging.Logger):
|
def __init__(self, model_path: str, auth_token: str, logger: logging.Logger):
|
||||||
super().__init__(model_path, auth_token, logger)
|
super().__init__(model_path, auth_token, logger)
|
||||||
self._model_path = model_path
|
self._model_path = model_path
|
||||||
|
@ -133,6 +133,12 @@ class HFTokenizer(HFHubBase):
|
||||||
with path.read_text(encoding='utf-8') as file:
|
with path.read_text(encoding='utf-8') as file:
|
||||||
return json.loads(file)
|
return json.loads(file)
|
||||||
|
|
||||||
|
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]:
|
def tokenizer_config(self, model_repo: str) -> dict[str, object]:
|
||||||
path = self.model_path / model_repo / "tokenizer_config.json"
|
path = self.model_path / model_repo / "tokenizer_config.json"
|
||||||
with path.read_text(encoding='utf-8') as file:
|
with path.read_text(encoding='utf-8') as file:
|
||||||
|
@ -143,11 +149,21 @@ class HFTokenizer(HFHubBase):
|
||||||
with path.read_text(encoding='utf-8') as file:
|
with path.read_text(encoding='utf-8') as file:
|
||||||
return json.loads(file)
|
return json.loads(file)
|
||||||
|
|
||||||
def tokenizer_model(self, model_repo: str) -> SentencePieceProcessor:
|
def get_normalizer(self, model_repo: str) -> None | dict[str, object]:
|
||||||
path = self.model_path / model_repo / "tokenizer.model"
|
normalizer = self.tokenizer_json(model_repo).get("normalizer", dict())
|
||||||
processor = SentencePieceProcessor()
|
if normalizer:
|
||||||
processor.LoadFromFile(path.read_bytes())
|
self.logger.info(f"JSON:Normalizer: {json.dumps(normalizer, indent=2)}")
|
||||||
return processor
|
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", dict())
|
||||||
|
if pre_tokenizer:
|
||||||
|
self.logger.info(f"JSON:PreTokenizer: {json.dumps(pre_tokenizer, indent=2)}")
|
||||||
|
else:
|
||||||
|
self.logger.warn(f"WARN:PreTokenizer: {pre_tokenizer}")
|
||||||
|
return pre_tokenizer
|
||||||
|
|
||||||
def get_tokenizer_json_hash(self, model_repo: str) -> str:
|
def get_tokenizer_json_hash(self, model_repo: str) -> str:
|
||||||
tokenizer = self.tokenizer_json(model_repo)
|
tokenizer = self.tokenizer_json(model_repo)
|
||||||
|
@ -169,7 +185,7 @@ class HFTokenizer(HFHubBase):
|
||||||
self.logger.info(f"JSON:PreTokenizer: {k}: {json.dumps(v, indent=2)}")
|
self.logger.info(f"JSON:PreTokenizer: {k}: {json.dumps(v, indent=2)}")
|
||||||
|
|
||||||
|
|
||||||
class HFVocabRequest(HFHubBase):
|
class HFModel(HFHub):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
auth_token: str,
|
auth_token: str,
|
||||||
|
@ -178,6 +194,16 @@ class HFVocabRequest(HFHubBase):
|
||||||
):
|
):
|
||||||
super().__init__(model_path, auth_token, logger)
|
super().__init__(model_path, auth_token, logger)
|
||||||
self._tokenizer = HFTokenizer(model_path, auth_token, logger)
|
self._tokenizer = HFTokenizer(model_path, auth_token, logger)
|
||||||
|
login(auth_token) # NOTE: Required for using model_info
|
||||||
|
|
||||||
|
@property
|
||||||
|
def model_type(self) -> ModelFileType:
|
||||||
|
return ModelFileType
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_model_info(repo_id: str) -> list[str]:
|
||||||
|
# NOTE: Get repository metadata to extract remote filenames
|
||||||
|
return [x.rfilename for x in model_info(repo_id).siblings]
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def tokenizer(self) -> HFTokenizer:
|
def tokenizer(self) -> HFTokenizer:
|
||||||
|
@ -194,40 +220,9 @@ class HFVocabRequest(HFHubBase):
|
||||||
self.logger.info(f"Downloaded tokenizer {file_name} from {model_repo}")
|
self.logger.info(f"Downloaded tokenizer {file_name} from {model_repo}")
|
||||||
|
|
||||||
def get_all_vocab_files(self, model_repo: str, vocab_type: VocabType) -> None:
|
def get_all_vocab_files(self, model_repo: str, vocab_type: VocabType) -> None:
|
||||||
vocab_list = HFTokenizer.get_vocab_filenames(vocab_type)
|
vocab_list = self.tokenizer.get_vocab_filenames(vocab_type)
|
||||||
for vocab_file in vocab_list:
|
for vocab_file in vocab_list:
|
||||||
dir_path = self.model_path / model_repo
|
dir_path = self.model_path / model_repo
|
||||||
file_path = dir_path / vocab_file
|
file_path = dir_path / vocab_file
|
||||||
os.makedirs(dir_path, exist_ok=True)
|
os.makedirs(dir_path, exist_ok=True)
|
||||||
self.get_vocab_file(model_repo, vocab_file, file_path)
|
self.get_vocab_file(model_repo, vocab_file, file_path)
|
||||||
|
|
||||||
def get_normalizer(self, model_repo: str) -> None | dict[str, object]:
|
|
||||||
normalizer = self.tokenizer.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.tokenizer_json(model_repo).get("pre_tokenizer", dict())
|
|
||||||
if pre_tokenizer:
|
|
||||||
self.logger.info(f"JSON:PreTokenizer: {json.dumps(pre_tokenizer, indent=2)}")
|
|
||||||
else:
|
|
||||||
self.logger.warn(f"WARN:PreTokenizer: {pre_tokenizer}")
|
|
||||||
return pre_tokenizer
|
|
||||||
|
|
||||||
|
|
||||||
# TODO:
|
|
||||||
class HFModelRequest(HFHubBase):
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
model_path: None | str | pathlib.Path,
|
|
||||||
auth_token: str,
|
|
||||||
logger: None | logging.Logger
|
|
||||||
):
|
|
||||||
super().__init__(model_path, auth_token, logger)
|
|
||||||
|
|
||||||
@property
|
|
||||||
def model_type(self) -> ModelFileType:
|
|
||||||
return ModelFileType
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue