File size: 14,309 Bytes
4d96efa db148eb 4d96efa db148eb 4d96efa db148eb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 |
---
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
<img src="https://huggingface.co./umarbutler/emubert/raw/main/logo.png" width="100" align="left" />
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](https://huggingface.co./datasets/umarbutler/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](https://huggingface.co./umarbutler/emubert/blob/main/LICENCE.md).
## Usage 👩💻
Those interested in finetuning EmuBert can check out Hugging Face's documentation for [Roberta](https://huggingface.co./roberta-base)-like models [here](https://huggingface.co./docs/transformers/en/model_doc/roberta) 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](https://github.com/UKPLab/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:
```python
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](https://huggingface.co./datasets/umarbutler/open-australian-legal-corpus/tree/fe0cd918dbe0a1fb5afe09cfa682ec3dbc1b94ca) of the [Open Australian Legal Corpus](https://huggingface.co./datasets/umarbutler/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;
1. Return carriages followed by newlines were replaced with newlines;
1. Whitespace was removed from lines comprised entirely of whitespace;
1. Newlines and whitespace preceding newlines were removed from the end of texts;
1. Newlines and whitespace succeeding newlines were removed from the beginning of texts; and
1. 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](https://huggingface.co./roberta-base)-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](https://huggingface.co./roberta-base)'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](https://huggingface.co./docs/optimum/en/bettertransformer/overview) was used to enable fast path execution and scaled dot-product attention, alongside automatic mixed 16-bit precision and [bitsandbytes](https://huggingface.co./docs/bitsandbytes/main/en/reference/optim/adamw#bitsandbytes.optim.AdamW8bit)' 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](https://github.com/umarbutler/emubert-creator).
## Benchmarks 📊
EmuBert achieves a [(pseudo-)perplexity](https://doi.org/10.18653/v1/2020.acl-main.240) of 2.05 against [version 2.0.0](https://huggingface.co./datasets/umarbutler/open-australian-legal-qa/tree/b53a24f8edf5eb33d033a53b5b53d0a4a220d4ae) of the [Open Australian Legal QA](https://huggingface.co./datasets/umarbutler/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](https://huggingface.co./datasets/umarbutler/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](https://huggingface.co./umarbutler/emubert/blob/main/LICENCE.md).
## Citation 🔖
If you've relied on the model for your work, please cite:
```bibtex
@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](https://huggingface.co./datasets/umarbutler/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. |