Update app.py

app.py

@@ -9,10 +9,10 @@ import torchaudio.transforms as T
 
 import gradio as gr
 
-from transformers import AutoModelForCTC, Wav2Vec2BertProcessor
+from transformers import HubertForCTC, Wav2Vec2Processor
 
 # Config
-model_name = "Yehor/w2v-bert-2.0-uk-v2"
+model_name = "Yehor/mHuBERT-147-uk"
 
 min_duration = 0.5
 max_duration = 60
@@ -25,10 +25,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
 # Load the model
-asr_model = AutoModelForCTC.from_pretrained(model_name, torch_dtype=torch_dtype).to(
-    device
-)
-processor = Wav2Vec2BertProcessor.from_pretrained(model_name)
+asr_model = HubertForCTC.from_pretrained(model_name, torch_dtype=torch_dtype, device_map=device)
+processor = Wav2Vec2Processor.from_pretrained(model_name)
 
 if use_torch_compile:
     asr_model = torch.compile(asr_model)
@@ -156,13 +154,10 @@ def inference(audio_path, progress=gr.Progress()):
             resampler = T.Resample(sr, 16_000, dtype=audio_input.dtype)
             audio_input = resampler(audio_input)
 
-        audio_input = audio_input.squeeze().numpy()
-
-        features = processor([audio_input], sampling_rate=16_000).input_features
-        features = torch.tensor(features).to(device)
+        audio_input = audio_input.squeeze(0).numpy()
 
-        if torch_dtype == torch.float16:
-            features = features.half()
+        inputs = processor([audio_input], sampling_rate=16_000, padding=True).input_values
+        features = torch.tensor(np.array(inputs), dtype=torch_dtype).to(device)
 
         with torch.inference_mode():
             logits = asr_model(features).logits