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)