update readme

README.md

@@ -54,27 +54,21 @@ The model acheives the following mean average precision (mAP) when fine-tuned on
 You can extract its top-layer feature (and intermediate hidden states) using the following code:
 
 ```python
-import torch
-import torchaudio
-
 from transformers import AutoModel
-from datasets import load_dataset
-
-import pdb; pdb.set_trace()
-dataset = load_dataset("hf-internal-testing/librispeech_asr_demo", "clean", split="validation")
-sampling_rate = dataset.features["audio"].sampling_rate
-assert sampling_rate == 16000
+import torch
 
-device = torch.device("cpu")
+model = AutoModel.from_pretrained(
+    "marcoyang/spear-xlarge-speech-audio", 
+    trust_remote_code=True,
+    force_download=False,
+)
 if torch.cuda.is_available():
-    device = torch.device("cuda")
-
-model = AutoModel.from_pretrained("/mnt/shared-storage-user/housiyuan/xiaoyu/workspace/icefall_general_encoder/egs/general_audio_encoder/mtl/spear_large_speech_hf", trust_remote_code=True)
+    model = model.to("cuda")
 model.eval()
-model.to(device)
 
-audio = dataset[0]["audio"]["array"].to(device)
-audio_len = torch.tensor(audio.shape[-1]).to(device)
+device = next(model.parameters()).device
+audio = torch.randn(1, 160000).to(device) # dummy audio input of 10 seconds
+audio_len = torch.tensor([160000]).to(device)
 
 with torch.no_grad():
     outputs = model(audio, audio_len)
@@ -85,5 +79,7 @@ middle_out = outputs["hidden_states"] # list of (N,T,C)
 
 print(encoder_out)
 print(encoder_out_lens)
-print(middle_out[0].shape)
+print(len(middle_out)) # 13 layers
+print(middle_out[-1].shape)
+print(middle_out[-1])
 ```