optimum/legal-bert-base-uncased-neuronx

This repo contains artifacts from nlpaueb/legal-bert-base-uncased but in neuronx format compatible with INF2 and TRN1 devices.

from transformers import AutoTokenizer
from optimum.neuron import NeuronBertForMaskedLM

input_shapes = {"batch_size": 1, "sequence_length": 128}
compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
neuron_model = NeuronBertForMaskedLM.from_pretrained(
    "nlpaueb/legal-bert-base-uncased",
    export=True,
    **input_shapes,
    **compiler_args,
)
# Save locally
neuron_model.save_pretrained("legal_bert_masked_lm_neuronx")
neuron_model.push_to_hub(
    "legal_bert_masked_lm_neuronx",
    repository_id="optimum/legal-bert-base-uncased-neuronx",  # Replace with your HF Hub repo id
)