save_model.py

import argparse
from transformers import AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, AutoTokenizer


def push_model(model: PreTrainedModel, tokenizer: PreTrainedTokenizer, path: str) -> None:
    model.push_to_hub(path)
    tokenizer.push_to_hub(path)


def load_model(path: str) -> tuple[PreTrainedModel, PreTrainedTokenizer]:
    model = AutoModelForCausalLM.from_pretrained(
        path,
    )
    tokenizer = AutoTokenizer.from_pretrained(path)
    return model, tokenizer


def main():
    parser = argparse.ArgumentParser(description="Save and Load Model Script")
    parser.add_argument(
        "--model_path", type=str, required=True, help="Path to the model directory"
    )
    parser.add_argument(
        "--remote_path", type=str, required=True, help="Remote path to push the model"
    )
    args = parser.parse_args()

    model, tokenizer = load_model(args.model_path)
    print(f"Model loaded from {args.model_path}")

    print(
        "Do you want to proceed with uploading the pruned model to the Hugging Face Hub? (y/n): ",
        end="",
    )
    choice = input().strip().lower()
    if choice == "y":
        print("Uploading model to Hugging Face Hub...")
        push_model(model, tokenizer, args.remote_path)
        print("Upload complete.")
    else:
        print("Upload skipped.")


if __name__ == "__main__":
    main()