Spaces:
Sleeping
Sleeping
| import argparse | |
| import json | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| from typing import Optional, Tuple | |
| import torch | |
| try: | |
| from huggingface_hub import ( | |
| create_repo, | |
| get_hf_file_metadata, | |
| hf_hub_download, | |
| hf_hub_url, | |
| repo_type_and_id_from_hf_id, | |
| upload_folder, | |
| ) | |
| from huggingface_hub.utils import EntryNotFoundError | |
| _has_hf_hub = True | |
| except ImportError: | |
| _has_hf_hub = False | |
| from .factory import create_model_from_pretrained, get_model_config, get_tokenizer | |
| from .tokenizer import HFTokenizer | |
| def save_config_for_hf( | |
| model, | |
| config_path: str, | |
| model_config: Optional[dict] | |
| ): | |
| preprocess_cfg = { | |
| 'mean': model.visual.image_mean, | |
| 'std': model.visual.image_std, | |
| } | |
| hf_config = { | |
| 'model_cfg': model_config, | |
| 'preprocess_cfg': preprocess_cfg, | |
| } | |
| with config_path.open('w') as f: | |
| json.dump(hf_config, f, indent=2) | |
| def save_for_hf( | |
| model, | |
| tokenizer: HFTokenizer, | |
| model_config: dict, | |
| save_directory: str, | |
| weights_filename='open_clip_pytorch_model.bin', | |
| config_filename='open_clip_config.json', | |
| ): | |
| save_directory = Path(save_directory) | |
| save_directory.mkdir(exist_ok=True, parents=True) | |
| weights_path = save_directory / weights_filename | |
| torch.save(model.state_dict(), weights_path) | |
| tokenizer.save_pretrained(save_directory) | |
| config_path = save_directory / config_filename | |
| save_config_for_hf(model, config_path, model_config=model_config) | |
| def push_to_hf_hub( | |
| model, | |
| tokenizer, | |
| model_config: Optional[dict], | |
| repo_id: str, | |
| commit_message: str = 'Add model', | |
| token: Optional[str] = None, | |
| revision: Optional[str] = None, | |
| private: bool = False, | |
| create_pr: bool = False, | |
| model_card: Optional[dict] = None, | |
| ): | |
| if not isinstance(tokenizer, HFTokenizer): | |
| # default CLIP tokenizers use https://huggingface.co/openai/clip-vit-large-patch14 | |
| tokenizer = HFTokenizer('openai/clip-vit-large-patch14') | |
| # Create repo if it doesn't exist yet | |
| repo_url = create_repo(repo_id, token=token, private=private, exist_ok=True) | |
| # Infer complete repo_id from repo_url | |
| # Can be different from the input `repo_id` if repo_owner was implicit | |
| _, repo_owner, repo_name = repo_type_and_id_from_hf_id(repo_url) | |
| repo_id = f"{repo_owner}/{repo_name}" | |
| # Check if README file already exist in repo | |
| try: | |
| get_hf_file_metadata(hf_hub_url(repo_id=repo_id, filename="README.md", revision=revision)) | |
| has_readme = True | |
| except EntryNotFoundError: | |
| has_readme = False | |
| # Dump model and push to Hub | |
| with TemporaryDirectory() as tmpdir: | |
| # Save model weights and config. | |
| save_for_hf( | |
| model, | |
| tokenizer=tokenizer, | |
| model_config=model_config, | |
| save_directory=tmpdir, | |
| ) | |
| # Add readme if it does not exist | |
| if not has_readme: | |
| model_card = model_card or {} | |
| model_name = repo_id.split('/')[-1] | |
| readme_path = Path(tmpdir) / "README.md" | |
| readme_text = generate_readme(model_card, model_name) | |
| readme_path.write_text(readme_text) | |
| # Upload model and return | |
| return upload_folder( | |
| repo_id=repo_id, | |
| folder_path=tmpdir, | |
| revision=revision, | |
| create_pr=create_pr, | |
| commit_message=commit_message, | |
| ) | |
| def push_pretrained_to_hf_hub( | |
| model_name, | |
| pretrained: str, | |
| repo_id: str, | |
| image_mean: Optional[Tuple[float, ...]] = None, | |
| image_std: Optional[Tuple[float, ...]] = None, | |
| commit_message: str = 'Add model', | |
| token: Optional[str] = None, | |
| revision: Optional[str] = None, | |
| private: bool = False, | |
| create_pr: bool = False, | |
| model_card: Optional[dict] = None, | |
| ): | |
| model, preprocess_eval = create_model_from_pretrained( | |
| model_name, | |
| pretrained=pretrained, | |
| image_mean=image_mean, | |
| image_std=image_std, | |
| ) | |
| model_config = get_model_config(model_name) | |
| assert model_config | |
| tokenizer = get_tokenizer(model_name) | |
| push_to_hf_hub( | |
| model=model, | |
| tokenizer=tokenizer, | |
| model_config=model_config, | |
| repo_id=repo_id, | |
| commit_message=commit_message, | |
| token=token, | |
| revision=revision, | |
| private=private, | |
| create_pr=create_pr, | |
| model_card=model_card, | |
| ) | |
| def generate_readme(model_card: dict, model_name: str): | |
| readme_text = "---\n" | |
| readme_text += "tags:\n- zero-shot-image-classification\n- clip\n" | |
| readme_text += "library_tag: open_clip\n" | |
| readme_text += f"license: {model_card.get('license', 'mit')}\n" | |
| if 'details' in model_card and 'Dataset' in model_card['details']: | |
| readme_text += 'datasets:\n' | |
| readme_text += f"- {model_card['details']['Dataset'].lower()}\n" | |
| readme_text += "---\n" | |
| readme_text += f"# Model card for {model_name}\n" | |
| if 'description' in model_card: | |
| readme_text += f"\n{model_card['description']}\n" | |
| if 'details' in model_card: | |
| readme_text += f"\n## Model Details\n" | |
| for k, v in model_card['details'].items(): | |
| if isinstance(v, (list, tuple)): | |
| readme_text += f"- **{k}:**\n" | |
| for vi in v: | |
| readme_text += f" - {vi}\n" | |
| elif isinstance(v, dict): | |
| readme_text += f"- **{k}:**\n" | |
| for ki, vi in v.items(): | |
| readme_text += f" - {ki}: {vi}\n" | |
| else: | |
| readme_text += f"- **{k}:** {v}\n" | |
| if 'usage' in model_card: | |
| readme_text += f"\n## Model Usage\n" | |
| readme_text += model_card['usage'] | |
| readme_text += '\n' | |
| if 'comparison' in model_card: | |
| readme_text += f"\n## Model Comparison\n" | |
| readme_text += model_card['comparison'] | |
| readme_text += '\n' | |
| if 'citation' in model_card: | |
| readme_text += f"\n## Citation\n" | |
| if not isinstance(model_card['citation'], (list, tuple)): | |
| citations = [model_card['citation']] | |
| else: | |
| citations = model_card['citation'] | |
| for c in citations: | |
| readme_text += f"```bibtex\n{c}\n```\n" | |
| return readme_text | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Push to Hugging Face Hub") | |
| parser.add_argument( | |
| "--model", type=str, help="Name of the model to use.", | |
| ) | |
| parser.add_argument( | |
| "--pretrained", type=str, | |
| help="Use a pretrained CLIP model weights with the specified tag or file path.", | |
| ) | |
| parser.add_argument( | |
| "--repo-id", type=str, | |
| help="Destination HF Hub repo-id ie 'organization/model_id'.", | |
| ) | |
| parser.add_argument( | |
| '--image-mean', type=float, nargs='+', default=None, metavar='MEAN', | |
| help='Override default image mean value of dataset') | |
| parser.add_argument( | |
| '--image-std', type=float, nargs='+', default=None, metavar='STD', | |
| help='Override default image std deviation of of dataset') | |
| args = parser.parse_args() | |
| print(f'Saving model {args.model} with pretrained weights {args.pretrained} to Hugging Face Hub at {args.repo_id}') | |
| # FIXME add support to pass model_card json / template from file via cmd line | |
| push_pretrained_to_hf_hub( | |
| args.model, | |
| args.pretrained, | |
| args.repo_id, | |
| image_mean=args.image_mean, # override image mean/std if trained w/ non defaults | |
| image_std=args.image_std, | |
| ) | |
| print(f'{args.model} saved.') | |