Views
No views yet
1import torch
2from datasets import load_dataset
3from transformers import AutoProcessor, AutoFeatureExtractor
4
5from optimum.neuron import NeuronModelForXVector, pipeline
6
7
8dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
9dataset = dataset.sort("id")
10sampling_rate = dataset.features["audio"].sampling_rate
11
12model_id = "anton-l/wav2vec2-base-superb-sv"
13feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
14input_shapes = {"batch_size": 1, "audio_sequence_length": 100000}
15compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
16model = NeuronModelForXVector.from_pretrained(
17 model_id,
18 export=True,
19 disable_neuron_cache=True,
20 **input_shapes,
21 **compiler_args,
22)
23model.save_pretrained("wav2vec2_neuron")