| 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) | |