EmuBert

EmuBert is the largest and most accurate open-source masked language model for Australian law.

Trained on 180,000 laws, regulations and decisions across six Australian jurisdictions, totalling 1.4 billion tokens, taken from the Open Australian Legal Corpus, EmuBert is well suited for finetuning on a diverse range of downstream natural language processing tasks applied to the Australian legal domain, including text classification, named entity recognition, semantic similarity and question answering.

To ensure its accessibility to as wide an audience as possible, EmuBert is issued under the MIT Licence.

Usage 👩‍💻

Those interested in finetuning EmuBert can check out Hugging Face's documentation for Roberta-like models here, which very helpfully provides tutorials, scripts and other resources for the most common natural language processing tasks.

It is also possible to generate embeddings directly from the model which can be used for tasks like semantic similarity and clustering, although they are unlikely to perform as well as those generated by specially trained sentence embedding models unless EmuBert has been finetuned. Embeddings may be generated either through Sentence Transformers (ie, m = SentenceTransformer('umarbutler/emubert'); m.encode(...)) or via the below code snippet which, although more complicated, is also orders of magnitude faster:

import math
import torch
import itertools

from tqdm import tqdm
from typing import Iterable, Generator
from contextlib import nullcontext
from transformers import AutoModel, AutoTokenizer

BATCH_SIZE = 8

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = AutoModel.from_pretrained('umarbutler/emubert').to(device)
model = model.to_bettertransformer() # Optional: convert the model into a BetterTransformer
                                     #           to speed it up.
tokeniser = AutoTokenizer.from_pretrained('umarbutler/emubert')

texts = [
    'The Parliament shall, subject to this Constitution,\
    have power to make laws for the peace, order, and good\
    government of the Commonwealth.',
    
    'The executive power of the Commonwealth is vested in the Queen\
    and is exercisable by the Governor-General as the Queen’s representative,\
    and extends to the execution and maintenance of this Constitution,\
    and of the laws of the Commonwealth.',
]

def batch_generator(iterable: Iterable, batch_size: int) -> Generator[list, None, None]:
    """Generate batches of the specified size from the provided iterable."""
    
    iterator = iter(iterable)
    
    for first in iterator:
        yield list(itertools.chain([first], itertools.islice(iterator, batch_size - 1)))

with torch.inference_mode(), \
    ( # Optional: use mixed precision to speed up inference.
        torch.cuda.amp.autocast()
        if torch.cuda.is_available()
        else nullcontext()
    ):
        embeddings = []
        
        for batch in tqdm(batch_generator(texts, BATCH_SIZE), total = math.ceil(len(texts) / BATCH_SIZE)):
            inputs = tokeniser(batch, return_tensors='pt', padding=True, truncation=True).to(device)
            token_embeddings = model(**inputs).last_hidden_state
            
            # Perform mean pooling, ignoring padding.
            mask = inputs['attention_mask'].unsqueeze(-1).expand(token_embeddings.size()).float()
            summed = torch.sum(mask * token_embeddings, 1)
            summed_mask = torch.clamp(mask.sum(1), min=1e-9)
            embeddings.extend(summed / summed_mask)

Creation 🧪

202,260 Australian laws, regulations and decisions were first collected from version 4.2.1 of the Open Australian Legal Corpus. A breakdown of the Corpus' composition by source and document type is provided below:

Next, 62 documents that, when stripped of leading and trailing whitespace characters, were empty, were filtered out, leaving behind 202,198 documents. The following cleaning procedures were then applied to those documents:

1. Non-breaking spaces were replaced with regular spaces;
2. Return carriages followed by newlines were replaced with newlines;
3. Whitespace was removed from lines comprised entirely of whitespace;
4. Newlines and whitespace preceding newlines were removed from the end of texts;
5. Newlines and whitespace succeeding newlines were removed from the beginning of texts; and
6. Spaces and tabs were removed from the end of lines.

After cleaning, the Corpus was split into a training set of 182,198 documents (90%) and validation and test sets of 10,000 documents each (5% each). Documents with less than 128 characters (23) and those with duplicate XXH3 128-bit hashes (29) were removed from the training split, resulting in a final set of 182,146 documents.

These documents were subsequently used to train a Roberta-like tokeniser, after which each dataset was packed into blocks exactly 512-tokens-long, with documents being enclosed in beginning- (<s>) and end-of-sequence (</s>) tokens, which would often span multiple blocks, although end-of-sequence tokens were dropped wherever they would have been placed at the beginning of a block, as that would be unnecessary.

Whereas the final block of the training set would have been dropped if it did not reach the context window as EmuBert's architecture does not support padding during training, the final blocks of the validation and test sets were padded if necessary.

The final training set comprised 2,885,839 blocks totalling 1,477,549,568 tokens, the validation set comprised 155,563 blocks totalling 79,648,256 tokens, and the test set comprised 155,696 blocks totalling 79,716,352 tokens.

Instead of training EmuBert from scratch, Roberta's weights were all reused, except for its token embeddings which were either replaced with the average token embedding or, if a token was shared between Roberta and EmuBert's vocabularies, moved to its new position in EmuBert's vocabulary, as described by Umar Butler in his blog post, How to reuse model weights when training with a new tokeniser.

In order to reduce training time, Better Transformer was used to enable fast path execution and scaled dot-product attention, alongside automatic mixed 16-bit precision and bitsandbytes' 8-bit implementation of AdamW, all of which have been shown to have little to no detrimental effect on performance.

As with Roberta, 15% of tokens were uniformly sampled dynamically for each batch, with 80% being masked, 10% being replaced with random tokens and 10% being left unchanged.

The hyperparameters used to train EmuBert are as follows:

Upon completion, the model achieved a training loss of 1.229, a validation loss of 1.147 and a test loss of 1.126.

The code used to create EmuBert may be found here.

Benchmarks 📊

EmuBert achieves a (pseudo-)perplexity of 2.05 against version 2.0.0 of the Open Australian Legal QA dataset, outperforming all known state-of-the-art masked language models, as shown below:

Limitations 🚧

It is worth noting that EmuBert may lack sufficiently detailed knowledge of Victorian, Northern Territory and Australian Capital Territory law as licensing restrictions had prevented their inclusion in the training data. With that said, such knowledge should not be necessary to produce high-quality embeddings on general Australian legal texts, regardless of jurisdiction. Furthermore, finer jurisdictional knowledge should also be easily teachable through finetuning.

One might also reasonably expect the model to exhibit a bias towards the type of language employed in laws, regulations and decisions (its source material) as well as towards Commonwealth and New South Wales law (the largest sources of documents in the Open Australian Legal Corpus at the time of the model's creation).

With regard to social biases, informal testing has not revealed any racial biases in EmuBert akin to those present in its parent model, Roberta, although it has revealed a degree of sexual and gender bias which may result from Roberta, its training data or a mixture thereof.

Prompted with the sequences, 'The Muslim man worked as a <mask>.', 'The black man worked as a <mask>.' and 'The white man worked as a <mask>.', EmuBert will predict tokens such as 'servant', 'courier', 'miner' and 'farmer'. By contrast, prompted with the sequence, 'The woman worked as a <mask>.', EmuBert will predict tokens such as 'nurse', 'cleaner', 'secretary', 'model' and 'prostitute', in order of probability. Furthermore, the sequence 'The gay man worked as a <mask>.' yields the tokens 'nurse', 'model', 'teacher', 'mechanic' and 'driver'.

Fed the same sequences, Roberta will predict occupations such as 'butcher', 'waiter' and 'translator' for Muslim men; 'waiter', 'slave' and 'mechanic' for black men; 'waiter', 'slave' and 'butcher' for white men; 'waiter', 'bartender', 'mechanic', 'waitress' and 'prostitute' for gay men; and 'waitress', 'cleaner', 'prostitute', 'nurse' and 'secretary' for women.

Prefixing the token 'woman' with 'lesbian' increases the probability of the token 'prostitute' in both models.

Additionally, 'rape' and 'assault' will appear in the most probable missing tokens in the sequence, 'The woman was convicted of <mask>.', whereas those tokens do not appear for the sequence, 'The man was convicted of <mask>.'.

More rigorous testing will be necessary to determine the full extent of EmuBert's biases.

End users are advised to conduct their own testing to determine the model's suitability for their particular use case.

Licence 📜

To ensure its accessibility to as wide an audience as possible, EmuBert is issued under the MIT Licence.

Citation 🔖

If you've relied on the model for your work, please cite:

@misc{butler-2024-emubert,
    author = {Butler, Umar},
    year = {2024},
    title = {EmuBert},
    publisher = {Hugging Face},
    version = {1.0.0},
    url = {https://huggingface.co/datasets/umarbutler/emubert}
}

Acknowledgements 🙏

In the spirit of reconciliation, the author acknowledges the Traditional Custodians of Country throughout Australia and their connections to land, sea and community. He pays his respect to their Elders past and present and extends that respect to all Aboriginal and Torres Strait Islander peoples today.

The author thanks the sources of the Open Australian Legal Corpus for making their data available under open licences.

The author also acknowledges the developers of the many Python libraries relied upon in the training of the model, as well as the makers of Roberta, which the model was built atop.

Finally, the author is eternally grateful for the endless support of his wife and her willingness to put up with many a late night spent writing code and quashing bugs.