Views
No views yet
1
2from transformers import AutoTokenizer
3from optimum.neuron import NeuronBertForTokenClassification
4
5input_shapes = {"batch_size": 1, "sequence_length": 128}
6compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
7neuron_model = NeuronBertForTokenClassification.from_pretrained(
8 "dslim/bert-base-NER",
9 export=True,
10 **input_shapes,
11 **compiler_args,
12)
13# Save locally
14neuron_model.save_pretrained("bert_base_ner_neuronx")
15neuron_model.push_to_hub(
16 "bert_base_ner_neuronx",
17 repository_id="optimum/bert-base-NER-neuronx", # Replace with your HF Hub repo id
18)