updt example

README.md

@@ -130,23 +130,26 @@ import torchaudio
 
 from transformers import AutoModelForCTC, Wav2Vec2ProcessorWithLM
 
-model = AutoModelForCTC.from_pretrained("bhuang/asr-wav2vec2-french").cuda()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = AutoModelForCTC.from_pretrained("bhuang/asr-wav2vec2-french").to(device)
 processor_with_lm = Wav2Vec2ProcessorWithLM.from_pretrained("bhuang/asr-wav2vec2-french")
+model_sample_rate = processor_with_lm.feature_extractor.sampling_rate
 
 wav_path = "example.wav"  # path to your audio file
 waveform, sample_rate = torchaudio.load(wav_path)
 waveform = waveform.squeeze(axis=0)  # mono
 
 # resample
-if sample_rate != 16_000:
-    resampler = torchaudio.transforms.Resample(sample_rate, 16_000)
+if sample_rate != model_sample_rate:
+    resampler = torchaudio.transforms.Resample(sample_rate, model_sample_rate)
     waveform = resampler(waveform)
 
 # normalize
-input_dict = processor_with_lm(waveform, sampling_rate=16_000, return_tensors="pt")
+input_dict = processor_with_lm(waveform, sampling_rate=model_sample_rate, return_tensors="pt")
 
 with torch.inference_mode():
-    logits = model(input_dict.input_values.to("cuda")).logits
+    logits = model(input_dict.input_values.to(device)).logits
 
 predicted_sentence = processor_with_lm.batch_decode(logits.cpu().numpy()).text[0]
 ```
@@ -159,23 +162,26 @@ import torchaudio
 
 from transformers import AutoModelForCTC, Wav2Vec2Processor
 
-model = AutoModelForCTC.from_pretrained("bhuang/asr-wav2vec2-french").cuda()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = AutoModelForCTC.from_pretrained("bhuang/asr-wav2vec2-french").to(device)
 processor = Wav2Vec2Processor.from_pretrained("bhuang/asr-wav2vec2-french")
+model_sample_rate = processor.feature_extractor.sampling_rate
 
 wav_path = "example.wav"  # path to your audio file
 waveform, sample_rate = torchaudio.load(wav_path)
 waveform = waveform.squeeze(axis=0)  # mono
 
 # resample
-if sample_rate != 16_000:
-    resampler = torchaudio.transforms.Resample(sample_rate, 16_000)
+if sample_rate != model_sample_rate:
+    resampler = torchaudio.transforms.Resample(sample_rate, model_sample_rate)
     waveform = resampler(waveform)
 
 # normalize
-input_dict = processor(waveform, sampling_rate=16_000, return_tensors="pt")
+input_dict = processor(waveform, sampling_rate=model_sample_rate, return_tensors="pt")
 
 with torch.inference_mode():
-    logits = model(input_dict.input_values.to("cuda")).logits
+    logits = model(input_dict.input_values.to(device)).logits
 
 # decode
 predicted_ids = torch.argmax(logits, dim=-1)