def setUp(self) -> None:
if os.path.exists(REPO_NAME):
shutil.rmtree(REPO_NAME, onerror=set_write_permission_and_retry)
logger.info(f"Does {REPO_NAME} exist: {os.path.exists(REPO_NAME)}")
repo = Repository(
REPO_NAME,
clone_from=f"{USER}/{REPO_NAME}",
use_auth_token=self._token,
git_user="******",
git_email="*****@*****.**",
)
with repo.commit("Add file to main branch"):
with open("dummy_file.txt", "w+") as f:
f.write("v1")
self.first_commit_hash = repo.git_head_hash()
with repo.commit("Add file to main branch"):
with open("dummy_file.txt", "w+") as f:
f.write("v2")
with open("dummy_file_2.txt", "w+") as f:
f.write("v3")
self.second_commit_hash = repo.git_head_hash()
with repo.commit("Add file to other branch", branch="other"):
with open("dummy_file_2.txt", "w+") as f:
f.write("v4")
self.third_commit_hash = repo.git_head_hash()
def setUp(self) -> None:
repo = Repository(
REPO_NAME,
clone_from=f"{USER}/{REPO_NAME}",
use_auth_token=self._token,
git_user="******",
git_email="*****@*****.**",
)
with repo.commit("Add file to main branch"):
with open("dummy_file.txt", "w+") as f:
f.write("v1")
self.first_commit_hash = repo.git_head_hash()
with repo.commit("Add file to main branch"):
with open("dummy_file.txt", "w+") as f:
f.write("v2")
with open("dummy_file_2.txt", "w+") as f:
f.write("v3")
self.second_commit_hash = repo.git_head_hash()
with repo.commit("Add file to other branch", branch="other"):
with open("dummy_file_2.txt", "w+") as f:
f.write("v4")
self.third_commit_hash = repo.git_head_hash()
def push_to_hf_hub(
model,
local_dir,
repo_namespace_or_url=None,
commit_message='Add model',
use_auth_token=True,
git_email=None,
git_user=None,
revision=None,
model_config=None,
):
if repo_namespace_or_url:
repo_owner, repo_name = repo_namespace_or_url.rstrip('/').split(
'/')[-2:]
else:
if isinstance(use_auth_token, str):
token = use_auth_token
else:
token = HfFolder.get_token()
if token is None:
raise ValueError(
"You must login to the Hugging Face hub on this computer by typing `transformers-cli login` and "
"entering your credentials to use `use_auth_token=True`. Alternatively, you can pass your own "
"token as the `use_auth_token` argument.")
repo_owner = HfApi().whoami(token)['name']
repo_name = Path(local_dir).name
repo_url = f'https://huggingface.co/{repo_owner}/{repo_name}'
repo = Repository(
local_dir,
clone_from=repo_url,
use_auth_token=use_auth_token,
git_user=git_user,
git_email=git_email,
revision=revision,
)
# Prepare a default model card that includes the necessary tags to enable inference.
readme_text = f'---\ntags:\n- image-classification\n- timm\nlibrary_tag: timm\n---\n# Model card for {repo_name}'
with repo.commit(commit_message):
# Save model weights and config.
save_for_hf(model, repo.local_dir, model_config=model_config)
# Save a model card if it doesn't exist.
readme_path = Path(repo.local_dir) / 'README.md'
if not readme_path.exists():
readme_path.write_text(readme_text)
return repo.git_remote_url()
def push_to_hf_hub(model: Any, model_name: str, task: str, **kwargs) -> None:
"""Save model and its configuration on HF hub
>>> from doctr.models import login_to_hub, push_to_hf_hub
>>> from doctr.models.recognition import crnn_mobilenet_v3_small
>>> login_to_hub()
>>> model = crnn_mobilenet_v3_small(pretrained=True)
>>> push_to_hf_hub(model, 'my-model', 'recognition', arch='crnn_mobilenet_v3_small')
Args:
model: TF or PyTorch model to be saved
model_name: name of the model which is also the repository name
task: task name
**kwargs: keyword arguments for push_to_hf_hub
"""
run_config = kwargs.get("run_config", None)
arch = kwargs.get("arch", None)
if run_config is None and arch is None:
raise ValueError("run_config or arch must be specified")
if task not in [
"classification", "detection", "recognition", "obj_detection"
]:
raise ValueError(
"task must be one of classification, detection, recognition, obj_detection"
)
# default readme
readme = textwrap.dedent(f"""
---
language: en
---
<p align="center">
<img src="https://github.com/mindee/doctr/releases/download/v0.3.1/Logo_doctr.gif" width="60%">
</p>
**Optical Character Recognition made seamless & accessible to anyone, powered by TensorFlow 2 & PyTorch**
## Task: {task}
https://github.com/mindee/doctr
### Example usage:
```python
>>> from doctr.io import DocumentFile
>>> from doctr.models import ocr_predictor, from_hub
>>> img = DocumentFile.from_images(['<image_path>'])
>>> # Load your model from the hub
>>> model = from_hub('mindee/my-model')
>>> # Pass it to the predictor
>>> # If your model is a recognition model:
>>> predictor = ocr_predictor(det_arch='db_mobilenet_v3_large',
>>> reco_arch=model,
>>> pretrained=True)
>>> # If your model is a detection model:
>>> predictor = ocr_predictor(det_arch=model,
>>> reco_arch='crnn_mobilenet_v3_small',
>>> pretrained=True)
>>> # Get your predictions
>>> res = predictor(img)
```
""")
# add run configuration to readme if available
if run_config is not None:
arch = run_config.arch
readme += textwrap.dedent(f"""### Run Configuration
\n{json.dumps(vars(run_config), indent=2, ensure_ascii=False)}"""
)
if arch not in AVAILABLE_ARCHS[task]: # type: ignore
raise ValueError(f"Architecture: {arch} for task: {task} not found.\
\nAvailable architectures: {AVAILABLE_ARCHS}")
commit_message = f"Add {model_name} model"
local_cache_dir = os.path.join(os.path.expanduser("~"), ".cache",
"huggingface", "hub", model_name)
repo_url = HfApi().create_repo(model_name,
token=HfFolder.get_token(),
exist_ok=False)
repo = Repository(local_dir=local_cache_dir,
clone_from=repo_url,
use_auth_token=True)
with repo.commit(commit_message):
_save_model_and_config_for_hf_hub(model,
repo.local_dir,
arch=arch,
task=task)
readme_path = Path(repo.local_dir) / "README.md"
readme_path.write_text(readme)
repo.git_push()
