upload.py 1.32 KB
Newer Older
mashun1's avatar
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import argparse

from transformers import AutoModelForTokenClassification, AutoTokenizer


def main():
    parser = argparse.ArgumentParser(
        description="Upload a trained model and tokenizer to the Hugging Face Hub."
    )
    parser.add_argument(
        "--model_path",
        type=str,
        required=True,
        help="Local path to the saved model directory (contains model and tokenizer files).",
    )
    parser.add_argument(
        "--repo_id",
        type=str,
        required=True,
        help="Target repository id on Hugging Face (e.g., KRLabsOrg/lettucedect-base-modernbert-en-v1).",
    )
    parser.add_argument(
        "--use_auth_token",
        action="store_true",
        help="Include this flag to use your Hugging Face authentication token (if not already set up).",
    )
    args = parser.parse_args()

    print(f"Loading model and tokenizer from {args.model_path} ...")
    model = AutoModelForTokenClassification.from_pretrained(args.model_path)
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)

    print(f"Uploading model to Hugging Face Hub at repo: {args.repo_id} ...")
    model.push_to_hub(args.repo_id, use_auth_token=args.use_auth_token)
    tokenizer.push_to_hub(args.repo_id, use_auth_token=args.use_auth_token)
    print("Upload complete!")


if __name__ == "__main__":
    main()