Function push_to_hub

scripts/generate_tiny_models/_common.py:208–253

(model, tokenizer, generation_config, prefix=None, suffix=None, create_pr=None)

def push_to_hub(model, tokenizer, generation_config, prefix=None, suffix=None, create_pr=None):
    if create_pr is None:
        create_pr = _parse_args().create_pr

    model_class_name = model.__class__.__name__
    content = MODEL_CARD.format(model_class_name=model_class_name)
    model_card = ModelCard(content)
    if prefix is not None:
        model_class_name = f"{prefix}-{model_class_name}"
    repo_id = f"{ORGANIZATION}/{model_class_name}"
    if suffix is not None:
        repo_id += f"-{suffix}"

    exists = api.repo_exists(repo_id)
    if exists and not create_pr:
        print(f"Model {repo_id} already exists, skipping (pass --create-pr to open a PR)")
        return

    if not exists:
        api.create_repo(repo_id, exist_ok=True)

    # Save all artifacts to a temp dir and upload them in a single commit, so --create-pr opens one PR.
    with tempfile.TemporaryDirectory() as tmpdir:
        model.save_pretrained(tmpdir)
        if tokenizer is not None:
            tokenizer.save_pretrained(tmpdir)
        if generation_config is not None:
            generation_config.save_pretrained(tmpdir)
        model_card.save(os.path.join(tmpdir, "README.md"))

        operations = [
            CommitOperationAdd(
                path_in_repo=os.path.relpath(os.path.join(root, name), tmpdir),
                path_or_fileobj=os.path.join(root, name),
            )
            for root, _, files in os.walk(tmpdir)
            for name in files
        ]
        commit_info = api.create_commit(
            repo_id=repo_id,
            operations=operations,
            commit_message=f"Upload {model.__class__.__name__}",
            create_pr=exists and create_pr,
        )
        if commit_info.pr_url:
            print(f"[push_to_hub] PR opened: {commit_info.pr_url}")


def init_weights_tiny_model(model):