49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Necessary to load the local gguf package
|
|
if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / 'gguf-py').exists():
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from gguf.constants import MODEL_ARCH, MODEL_ARCH_NAMES
|
|
from gguf.huggingface_hub import HFVocabRequest
|
|
|
|
logger = logging.getLogger("gguf-gen-pre")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("hf_auth_token", help="A huggingface read auth token")
|
|
parser.add_argument(
|
|
"-v", "--verbose", action="store_true", help="Increase output verbosity."
|
|
)
|
|
parser.add_argument(
|
|
"-r", "--model-repo", default="meta-llama/Llama-2-7b-hf",
|
|
help="The models repository. Default is 'meta-llama/Llama-2-7b-hf'."
|
|
)
|
|
parser.add_argument(
|
|
"-m", "--model-path", default="models/",
|
|
help="The models storage path. Default is 'models/'."
|
|
)
|
|
parser.add_argument(
|
|
"--vocab-type",
|
|
const="BPE", nargs="?", choices=["BPE", "SPM"],
|
|
help="The type of vocab. Default is 'BPE'."
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
vocab_request = HFVocabRequest(args.auth_token, args.model_path, logger)
|
|
vocab_type = vocab_request.get_vocab_enum(args.vocab_type)
|
|
vocab_request.get_all_vocab_files(args.model_repo, vocab_type)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|