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