Show full text

app.py

@@ -8,7 +8,7 @@ model_name = "Yehor/wav2vec2-xls-r-1b-uk-with-lm"
 tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(model_name)
 processor = Wav2Vec2ProcessorWithLM.from_pretrained(model_name)
 model = Wav2Vec2ForCTC.from_pretrained(model_name)
-model.to("cpu")
+model.to("cuda")
 
 
 # define function to read in sound file
@@ -45,6 +45,7 @@ def inference(audio):
         stride_length_s=(4, 2),
         return_tensors="pt",
     ).input_values
+    input_values = input_values.cuda()
 
     with torch.no_grad():
         logits = model(input_values).logits
@@ -55,6 +56,7 @@ def inference(audio):
     time_offset = 320 / sample_rate
 
     total_prediction = []
+    words = []
     for item in prediction.word_offsets:
         r = item
 
@@ -62,8 +64,11 @@ def inference(audio):
         e = round(r['end_offset'] * time_offset, 2)
 
         total_prediction.append(f"{s} - {e}: {r['word']}")
+        words.append(r['word'])
+
     print(prediction[0])
-    return "\n".join(total_prediction)
+
+    return "\n".join(total_prediction) + "\n\n" + ' '.join(words)
 
 
 inputs = gr.inputs.Audio(label="Input Audio", type="file")