File size: 641 Bytes
4baaf20
 
 
41a0b6d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
---
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
)
```