| | --- |
| | license: mit |
| | --- |
| | |
| | **INFERENTIA ONLY** |
| |
|
| | ```py |
| | |
| | from transformers import AutoTokenizer |
| | from optimum.neuron import NeuronBertForTokenClassification |
| | |
| | input_shapes = {"batch_size": 1, "sequence_length": 128} |
| | compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"} |
| | neuron_model = NeuronBertForTokenClassification.from_pretrained( |
| | "dslim/bert-base-NER", |
| | export=True, |
| | **input_shapes, |
| | **compiler_args, |
| | ) |
| | # Save locally |
| | neuron_model.save_pretrained("bert_base_ner_neuronx") |
| | neuron_model.push_to_hub( |
| | "bert_base_ner_neuronx", |
| | repository_id="optimum/bert-base-NER-neuronx", # Replace with your HF Hub repo id |
| | ) |
| | ``` |