Update README.md

README.md

@@ -59,20 +59,34 @@ from transformers import Wav2Vec2Processor
 from transformers import AutoModel
 import torch
 from torch import nn
+import torchaudio.transforms as T
 from datasets import load_dataset
 
+
 # load demo audio and set processor
 dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
 dataset = dataset.sort("id")
 sampling_rate = dataset.features["audio"].sampling_rate
-processor = Wav2Vec2Processor.from_pretrained("facebook/hubert-large-ls960-ft")
+processor = Wav2Vec2Processor.from_pretrained("m-a-p/MERT-v1-330M")
+
+resample_rate = processor.feature_extractor.sampling_rate
+# make sure the sample_rate aligned
+if resample_rate != sampling_rate:
+  resampler = T.Resample(sample_rate, resample_rate)
+else:
+  resampler = None
 
 # loading our model weights
 commit_hash='7bab7bb5d8b52448eff4873a980dc17f0015a09c'# this is recommended for security reason, the hash might be updated
 model = AutoModel.from_pretrained("m-a-p/MERT-v1-330M", trust_remote_code=True, revision=commit_hash)
 
 # audio file is decoded on the fly
-inputs = processor(dataset[0]["audio"]["array"], sampling_rate=sampling_rate, return_tensors="pt")
+if resampler is None:
+	input_audio = dataset[0]["audio"]["array"]
+else:
+  input_audio = resampler(dataset[0]["audio"]["array"])
+  
+inputs = processor(input_audio, sampling_rate=resample_rate, return_tensors="pt")
 with torch.no_grad():
     outputs = model(**inputs, output_hidden_states=True)