README.md

metadata

language:
  - en
license: mit
library_name: transformers
base_model: roberta-base
tags:
  - law
  - legal
  - australia
  - generated_from_trainer
  - fill-mask
  - sentence-similarity
  - feature-extraction
datasets:
  - umarbutler/open-australian-legal-corpus
widget:
  - text: >-
      Section <mask> of the Constitution grants the Australian Parliament the
      power to make laws for the peace, order, and good government of the
      Commonwealth.
  - text: The most learned and eminent jurist in Australia's history is <mask> CJ.
  - text: >-
      A <mask> of trade is valid to the extent to which it is not against public
      policy, whether it is in severable terms or not.
  - text: >-
      In Mabo v <mask> (No 2) (1992) 175 CLR 1, the Court found that the
      doctrine of terra nullius was not applicable to Australia at the time of
      British settlement of New South Wales.
metrics:
  - perplexity
model-index:
  - name: emubert
    results:
      - task:
          type: fill-mask
          name: Fill mask
        dataset:
          type: umarbutler/open-australian-legal-qa
          name: Open Australian Legal QA
          split: train
          revision: b53a24f8edf5eb33d033a53b5b53d0a4a220d4ae
        metrics:
          - type: perplexity
            value: 2.05
            name: perplexity
        source:
          name: EmuBert Creator
          url: https://github.com/umarbutler/emubert-creator

EmuBert

EmuBert is the largest 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 a diverse range of downstream natural language processing tasks applied to the Australian legal domain, including text classification, named entity recognition and question answering. It can also be used as-is for semantic similarity, vector search and general sentence embedding.

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 from the model which can be directly used for tasks like semantic similarity and clustering or for the training of downstream models. This can be done 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 torch
import itertools

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('model').to(device)
model = model.to_bettertransformer() # Optional: convert the model into a BetterTransformer
                                     #           to speed it up.
tokeniser = AutoTokenizer.from_pretrained('model')

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 batch_generator(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:

Source	Primary Legislation	Secondary Legislation	Bills	Decisions	Total
Federal Register of Legislation	3,872	19,587	0	0	23,459
Federal Court of Australia	0	0	0	46,733	46,733
High Court of Australia	0	0	0	9,433	9,433
NSW Caselaw	0	0	0	111,882	111,882
NSW Legislation	1,428	800	0	0	2,228
Queensland Legislation	564	426	2,247	0	3,237
Western Australian Legislation	812	760	0	0	1,572
South Australian Legislation	557	471	154	0	1,182
Tasmanian Legislation	858	1,676	0	0	2,534
Total	8,091	23,720	2,401	168,048	202,260

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.

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:

Hyperparameter	EmuBert	Roberta
Optimiser	AdamW 8-bit	Adam
Scheduler	Cosine	Linear
Precision	16-bit	16-bit
Batch size	8	8,000
Steps	1,000,000	500,000
Warmup steps	48,000	24,000
Learning rate	1e-5	6e-4
Weight decay	0.01	0.01
Adam epsilon	1e-6	1e-6
Adam beta1	0.9	0.9
Adam beta2	0.98	0.98
Gradient clipping	1	0

Upon completion, the model achieved a training loss of 1.229, 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:

Model	Perplexity
EmuBert	2.05
Roberta	2.38
Albert v2	3.49
Bert (cased)	2.18
Bert (uncased)	2.41

Limitations 🚧

Although informal testing has not revealed any racial, sexual, gender or other social biases, given that Roberta's weights were reused, it is possible that there may be some biases that have been transferred over to EmuBert. It is also possible that there are social biases present in the Corpus that may have been introduced via training. More rigorous testing is necessary to determine the true extent of any biases present in EmuBert.

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

Finally, it is worth noting that the model may lack 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. Furthermore, such knowledge should be easily teachable through finetuning.

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.