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