File size: 682 Bytes
1a0834b | 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 | ---
license: apache-2.0
---
**INFERENTIA2 ONLY**
```py
from transformers import AutoTokenizer
from optimum.neuron import NeuronBertForQuestionAnswering
input_shapes = {"batch_size": 1, "sequence_length": 128}
compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
neuron_model = NeuronBertForQuestionAnswering.from_pretrained(
"deepset/bert-base-cased-squad2",
export=True,
**input_shapes,
**compiler_args,
)
# Save locally
neuron_model.save_pretrained("bert_base_cased_squad2_neuronx")
neuron_model.push_to_hub(
"bert_base_cased_squad2_neuronx",
repository_id="optimum/bert-base-cased-squad2-neuronx", # Replace with your HF Hub repo id
)
``` |