transformer-smaller-training-vocab

import torch from transformers import AutoTokenizer, AutoModelForSequenceClassification from transformer_smaller_training_vocab import reduce_train_vocab_and_context, recreate_vocab # 1. Load a pre-trained model and tokenizer model_name = "bert-base-uncased" # Using a small, common model for quickstart tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForSequenceClassification.from_pretrained(model_name) # Dummy dataset for demonstration texts = ["hello world", "this is a test", "another example sentence"] tokenized_texts = [tokenizer(text, return_tensors="pt") for text in texts] # 2. Reduce the vocabulary reduced_tokenizer, reduced_model, added_tokens_during_reduction = reduce_train_vocab_and_context( model=model, tokenizer=tokenizer, tokenized_datasets=[t["input_ids"] for t in tokenized_texts], model_resize_strategy="embedding_resize", ) print(f"Original vocab size: {tokenizer.vocab_size}") print(f"Reduced vocab size: {reduced_tokenizer.vocab_size}") # You would typically train with reduced_tokenizer and reduced_model here # 3. Recreate the original vocabulary (after training, if needed) recreated_model = recreate_vocab( reduced_model=reduced_model, reduced_tokenizer=reduced_tokenizer, orig_tokenizer=tokenizer, added_tokens_during_reduction=added_tokens_during_reduction, model_resize_strategy="embedding_resize", ) print(f"Recreated model vocab size: {recreated_model.config.vocab_size}") # Verify recreated_model.config.vocab_size == tokenizer.vocab_size