feat: Add download_model method and fix references for clarity to mitigate confusion

This commit is contained in:
teleprint-me 2024-05-17 23:00:12 -04:00
parent 4790f76740
commit 5c8144e645
No known key found for this signature in database
GPG key ID: B0D11345E65C4D48

View file

@ -48,12 +48,11 @@ class HuggingFaceHub:
f.write(content) f.write(content)
self.logger.info(f"Wrote {len(content)} bytes to {path} successfully") self.logger.info(f"Wrote {len(content)} bytes to {path} successfully")
def resolve_path(self, repo: str, file: str) -> str: def resolve_url(self, repo: str, file: str) -> str:
return f"{self._base_url}/{repo}/resolve/main/{file}" return f"{self._base_url}/{repo}/resolve/main/{file}"
def download_file(self, repo: str, file: str): def download_file(self, url: str) -> requests.Response:
resolve_path = self.resolve_path(repo, file) response = self._session.get(url, headers=self.headers)
response = self._session.get(resolve_path, headers=self.headers)
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
@ -62,19 +61,22 @@ class HuggingFaceHub:
class HFTokenizerRequest: class HFTokenizerRequest:
def __init__( def __init__(
self, self,
dl_path: pathlib.Path, dl_path: None | str | pathlib.Path,
auth_token: str, auth_token: str,
logger: None | logging.Logger logger: None | logging.Logger
): ):
self._hub = HuggingFaceHub(auth_token, logger) self._hub = HuggingFaceHub(auth_token, logger)
self._models = MODEL_REPOS
if dl_path is None: if dl_path is None:
self._download_path = pathlib.Path("models/tokenizers") self._local_path = pathlib.Path("models/tokenizers")
else: else:
self._download_path = dl_path self._local_path = dl_path
self._files = ["config.json", "tokenizer_config.json", "tokenizer.json"] # Set the logger
if logger is None:
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger("hf-tok-req")
self.logger = logger
@property @property
def hub(self) -> HuggingFaceHub: def hub(self) -> HuggingFaceHub:
@ -82,21 +84,37 @@ class HFTokenizerRequest:
@property @property
def models(self) -> list[dict[str, object]]: def models(self) -> list[dict[str, object]]:
return self._models return MODEL_REPOS
@property @property
def download_path(self) -> pathlib.Path: def tokenizer_type(self) -> TokenizerType:
return self._download_path return TokenizerType
@download_path.setter
def download_path(self, value: pathlib.Path):
self._download_path = value
@property @property
def files(self) -> list[str]: def local_path(self) -> pathlib.Path:
return self._files return self._local_path
def download_file_with_auth(self, repo, file, directory): @local_path.setter
response = self.hub.download_file(repo, file) def local_path(self, value: pathlib.Path):
os.makedirs(os.path.dirname(directory), exist_ok=True) self._local_path = value
self.hub.write_file(response.content, directory)
def download_model(self) -> None:
for model in self.models:
name, repo, tokt = model['name'], model['repo'], model['tokt']
os.makedirs(f"{self.local_path}/{name}", exist_ok=True)
filenames = ["config.json", "tokenizer_config.json", "tokenizer.json"]
if tokt == self.tokenizer_type.SPM:
filenames.append("tokenizer.model")
for file_name in filenames:
file_path = pathlib.Path(f"{self.local_path}/{name}/{file_name}")
if file_path.is_file():
self.logger.info(f"skipped pre-existing tokenizer {name} at {file_path}")
continue
try: # NOTE: Do not use bare exceptions! They mask issues!
resolve_url = self.hub.resolve_url(repo, file_name)
response = self.hub.download_file(resolve_url)
self.hub.write_file(response.content, file_path)
except requests.exceptions.HTTPError as e:
self.logger.error(f"Failed to download tokenizer {name}: {e}")