File size: 891 Bytes
00a518f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import argparse
from transformers import AutoConfig, AutoTokenizer
def save_model_assets(model_name, output_dir):
# Load the configuration
config = AutoConfig.from_pretrained(model_name)
config.save_pretrained(output_dir)
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.save_pretrained(output_dir)
print(f"Configuration and tokenizer saved to {output_dir}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Save model config and tokenizer locally")
parser.add_argument("--model_name", type=str, required=True, help="Name of the model to load (e.g., NbAiLab/nb-bert-large)")
parser.add_argument("--output_dir", type=str, required=True, help="Directory to save the config and tokenizer")
args = parser.parse_args()
save_model_assets(args.model_name, args.output_dir)
|