|
--- |
|
license: apache-2.0 |
|
--- |
|
|
|
This is model is compiled explicitly for AWS Neuronx(inferentia 2 / trainium 1) with the following codes: |
|
|
|
```python |
|
from datasets import load_dataset |
|
from transformers import AutoProcessor |
|
|
|
from optimum.neuron import NeuronModelForAudioClassification, pipeline |
|
|
|
|
|
dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation") |
|
dataset = dataset.sort("id") |
|
sampling_rate = dataset.features["audio"].sampling_rate |
|
|
|
model_id = "facebook/wav2vec2-large-960h-lv60-self" |
|
processor = AutoProcessor.from_pretrained(model_id) |
|
input_shapes = {"batch_size": 1, "audio_sequence_length": 100000} |
|
compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"} |
|
model = NeuronModelForAudioClassification.from_pretrained( |
|
model_id, |
|
export=True, |
|
disable_neuron_cache=True, |
|
**input_shapes, |
|
**compiler_args, |
|
) |
|
model.save_pretrained("wav2vec2_neuron") |
|
``` |
