Building Your Own Tokenizer for Medical and Legal Documents
Today I want to talk about a part of the LLM pipeline that most people never check, but it breaks a lot of domain specific apps: The Tokenizer.If you have used a model like GPT or LLaMA on medical or legal text, you may have noticed something feels off. Drug names come out wrong. Legal citations…
Today I want to talk about a part of the LLM pipeline that most people never check, but it breaks a lot of domain specific apps: The Tokenizer.If you have used a model like GPT or LLaMA on medical or legal text, you may have noticed something feels off. Drug names come out wrong. Legal citations get slightly changed. If a model wasn’t trained on information about a rare condition, it won’t have that knowledge. In such cases, prompting alone can’t make the model accurately answer questions about it.The Real Problem: When Tokenizers Fail On Domain TextSay you are building a clinical notes assistant for a hospital. A doctor writes “Prescribed 12.5 mg hydrochlorothiazide daily.” You would expect maybe fifteen tokens.Instead, a general tokenizer breaks the word “hydrochlorothiazide” into five small pieces, because this word barely shows up in the normal web text the tokenizer was trained on.Now think about legal text. A contract mentions “42 U.S.C. § 1395.” A normal tokenizer has no idea this is one single legal reference. It splits it into ten or more pieces. This means there is a real chance the model gets the number slightly wrong later. In a legal document, that is not a small mistake.This is exactly the kind of problem a custom tokenizer fixes.In this article, I will show you why this happens and how to build a tokenizer that actually understands medical and legal text.What Is Tokenizer Fertility?There is a simple number that tells you how badly a tokenizer is cutting up your text. It is called fertility: the average number of tokens used per word.A few things to know about fertility:Under 1.5 is good. This means the tokenizer matches your text well.Above 2.0 is a warning sign. Most general tokenizers score this high on medical and legal text.It decides your real context window. A fertility of 2.0 means a “128K context” model really only gives you about 64K worth of your content.It decides your cost. More tokens per document means more money per document, every time.Once you check fertility on your own text, a lot of problems people blame on “the model just isn’t good with our data” turn out to be a tokenizer problem instead.Why Build A Custom Tokenizer Instead Of Just Fine Tuning?This is the question I get asked the most, so let me explain it simply.1. Fine Tuning Does Not Fix The Cutting ProblemFine tuning teaches the model to work better with the broken pieces it already has. It does not stop the breaking from happening in the first place. You still pay extra on every request, and the model is still trying to guess “hydrochlorothiazide” from five small pieces, just a little better than before.2. The Cost Adds Up FastIf you process two million documents a day and your tokenizer uses twice as many tokens as it should, you are not paying a little extra. You are paying for an entire second copy of tokens, every single day, forever.3. Safety MattersDrug dosages, legal citations, and medical codes should never be left to chance. A custom tokenizer lets you make sure these stay exactly correct, every time.4. You Get Your Full Context Window BackA model that says “128K context” only really gives you that much space if the tokenizer matches your text well. If not, you lose half of it to wasted, broken up tokens without even knowing it.The Building Blocks: Three Ways To TokenizeBefore we build anything, let’s look at the three common methods.For medical and legal text, I recommend BPE. Not because it is the smartest option, but because it fits. Almost every model you will want to build on top of already uses BPE, so your new medical or legal words slot right in instead of forcing you to rebuild everything from scratch.Building It Step By StepStep 1: Collect Clean Domain TextYour tokenizer is only as good as the text you train it on. For medical text, use PubMed articles for broad coverage, plus real (de-identified) clinical notes, since doctors write very differently than research papers. For legal text, use public sources like the Caselaw Access Project and CourtListener, plus a firm’s own contracts for real world language.Removing personal information here is not optional. If a patient’s name or ID shows up often enough in your text, it can actually become its own token, which means it gets baked into the tokenizer itself. Clean the data first, then check a sample by hand, before you train anything on it.Step 2: Train The Base Vocabularyfrom tokenizers import Tokenizer, models, trainers, pre_tokenizers, decodersdef train_bpe_tokenizer(corpus_path: str, vocab_size: int = 32000) -> Tokenizer: tokenizer = Tokenizer(models.BPE(unk_token="")) tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) tokenizer.decoder = decoders.ByteLevel() trainer = trainers.BpeTrainer( vocab_size=vocab_size, min_frequency=5, special_tokens=["", "", "", ""] ) tokenizer.train(files=[corpus_path], trainer=trainer) return tokenizercustom_tokenizer = train_bpe_tokenizer("medical_legal_corpus.txt", vocab_size=32000)custom_tokenizer.save("domain_tokenizer.json")Around 32,000 new words is a good starting point for either domain. Go too small, and drug names and citations keep breaking apart. Go too big, and you are paying extra cost for words that barely show up.Do not just trust that this helped. Check it. Here is a small script you can run right after training that measures fertility on a real sentence, side by side against a general tokenizer:def fertility(text: str, token_count: int) -> float: word_count = len(text.split()) return token_count / word_count if word_count else 0.0test_sentence = "Prescribed 12.5 mg hydrochlorothiazide daily for hypertension management."custom_count = len(custom_tokenizer.encode(test_sentence).ids)print(f"Custom tokenizer: {custom_count} tokens, fertility {fertility(test_sentence, custom_count):.2f}")# Compare against a general tokenizer, cl100k_base is the family GPT-4 usesimport tiktokenbaseline = tiktoken.get_encoding("cl100k_base")baseline_count = len(baseline.encode(test_sentence))print(f"Baseline tokenizer: {baseline_count} tokens, fertility {fertility(test_sentence, baseline_count):.2f}")I ran this exact setup on a small sample corpus while writing this article. The custom tokenizer took an eleven word sentence with “hydrochlorothiazide” in it down to 11 tokens, fertility 1.38.Custom tokenizer: 11 tokens, fertility 1.38Baseline tokenizer: 17 tokens, fertility 2.12A generic tokenizer trained without domain text needed 53 tokens for the same sentence, fertility >2. Your real numbers will differ based on your corpus size, but this is the exact check to run before you trust any of this is helping.Step 3: Protect Important TermsSome words are too important to leave to chance, even in a well trained tokenizer. A rare dosage or an unusual citation might not get its own token just from frequency alone. So instead of hoping for the best, we protect these directly with rules that run before the tokenizer even sees the text.import rePROTECTED_PATTERNS = [ re.compile(r"\b\d+(\.\d+)?\s?(mg|mcg|mL|units?)\b"), # dosages re.compile(r"\b[A-Z]\d{2}\.\d{1,2}\b"), # ICD-10 codes re.compile(r"\b\d{5}(-\d{2})?\b"), # CPT codes re.compile(r"\b\d+\s?U\.S\.C\.\s?§\s?\d+\b"), # US Code citations]def protect_terms(text: str) -> tuple[str, dict]: placeholders = {} for i, pattern in enumerate(PROTECTED_PATTERNS): for match in pattern.finditer(text): token_id = f"__SAVE_{i}_{len(placeholders)}__" placeholders[token_id] = match.group(0) text = text.replace(match.group(0), token_id, 1) return text, placeholdersRun this before tokenizing. Swap the placeholders back to the real text after. This way, a dosage or citation is always kept exactly correct, instead of hoping the tokenizer treats it well.Step 4: Add To An Existing Model Instead Of Starting OverYou do not need to train a whole new model to use your new tokenizer. Add your new medical or legal words on top of an existing model’s vocabulary, resize the embedding table, and set the starting values for the new words as the average of their old broken up pieces.Which words go in this list is not a guess. It is whatever your Step 2 fertility check just showed you was badly broken up. If “hydrochlorothiazide” needed 6 pieces under the old tokenizer, that is exactly the kind of word that belongs here.import torchdef extend_embeddings(model, old_vocab_size, new_tokens, old_tokenizer): new_vocab_size = old_vocab_size + len(new_tokens) old_embeddings = model.get_input_embeddings().weight.data hidden_dim = old_embeddings.shape[1] new_embeddings = torch.zeros(new_vocab_size, hidden_dim) new_embeddings[:old_vocab_size] = old_embeddings for i, token in enumerate(new_tokens): old_subword_ids = old_tokenizer.encode(token).ids new_embeddings[old_vocab_size + i] = old_embeddings[old_subword_ids].mean(dim=0) model.resize_token_embeddings(new_vocab_size) model.get_input_embeddings().weight.data = new_embeddings return model# The words to add come straight out of what Step 2 flagged as badly fragmented,# not a random guess.new_tokens = ["hydrochlorothiazide", "azithromycin", "res_judicata"]model = extend_embeddings(model, old_vocab_size=model.get_input_embeddings().weight.shape[0], new_tokens=new_tokens, old_tokenizer=old_tokenizer)One step people often skip: after adding the new words, you still need to train the model a bit more so it learns to actually use them. Skip this, and your new words start out looking fine on paper but the model has no real idea how to use them yet. This usually ends up worse than not adding them at all.Detailed Token Breakdown of a test sample :Here, in Baseline Tokenizer common words and especially the medical term hydrochlorothiazide are broken into multiple subword tokens because the tokenizer has never learned them as a single unit. With a domain-specific tokenizer, medical terms are treated as individual tokens instead of being fragmented into multiple pieces. This results in fewer tokens, preserves the semantic meaning of specialized terminology, and allows the model to process domain-specific text more efficiently.ConclusionA tokenizer that does not match your text is one of the most expensive, hardest to notice problems in a domain specific AI system. It quietly raises your cost, shrinks your usable context, and puts important details like dosages and legal citations at risk, all while the model still gives you an answer that looks fine on the surface. Fix the tokenizer, and you often get a bigger improvement than another round of prompt tweaking ever gives you.Thank you for reading! If this changed how you think about where your LLM costs are really going, give it a clap 👏 and share it with someone still trying to fix domain accuracy with prompts alone.This story is published on Generative AI. Connect with us on LinkedIn and follow Zeniteq to stay in the loop with the latest AI stories.Subscribe to our newsletter and YouTube channel to stay updated with the latest news and updates on generative AI. Let’s shape the future of AI together!Building Your Own Tokenizer for Medical and Legal Documents was originally published in Generative AI on Medium, where people are continuing the conversation by highlighting and responding to this story.Source: Generative AI Pub — Published — Category: Image AI