Used Nemo streaming logic

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import librosa
 import torch
 
-from math import floor,ceil
+from math import ceil
 import nemo.collections.asr as nemo_asr
 
 
@@ -17,11 +17,17 @@ asr_model.encoder.freeze()
 asr_model.decoder.freeze()
 
 
-total_buffer = asr_model.cfg["sample_rate"] * 60 // 10
-overhead_len = total_buffer // 2
-model_stride = 4
+buffer_len = 8.0
+chunk_len = 4.8
+total_buffer = round(buffer_len * asr_model.cfg.sample_rate)
+overhead_len = round((buffer_len - chunk_len) *  asr_model.cfg.sample_rate)
+model_stride = 8
 
 
+model_stride_in_secs = asr_model.cfg.preprocessor.window_stride * model_stride
+tokens_per_chunk = ceil(chunk_len / model_stride_in_secs)
+mid_delay = ceil((chunk_len + (buffer_len - chunk_len) / 2) / model_stride_in_secs)
+
 
 def resample(audio):
     audio_16k, sr = librosa.load(audio, sr = asr_model.cfg["sample_rate"], 
@@ -38,19 +44,13 @@ def model(audio_16k):
 
 
 def decode_predictions(logits_list):
-    # calc overhead
-    logits_overhead = logits_list[0].shape[1] * overhead_len / total_buffer / 2
-    if (logits_overhead * 2 != int(logits_overhead * 2)) and (len(logits_list) != 1):# if first chunk
-        raise ValueError("Wrong total_buffer")
-
+    logits_len = logits_list[0].shape[1]
     # cut overhead
     cutted_logits = []
     for idx in range(len(logits_list)):
-        start_cut = 0 if (idx==0) else floor(logits_overhead)
-        end_cut = 1 if (idx==len(logits_list)-1) else ceil(logits_overhead)
-        if (logits_overhead == int(logits_overhead)) and (end_cut != 1):
-            end_cut +=1
-        logits = logits_list[idx][:, start_cut:-end_cut]
+        start_cut = 0 if (idx==0) else logits_len - 1 - mid_delay
+        end_cut = -1 if (idx==len(logits_list)-1) else logits_len - 1 - mid_delay + tokens_per_chunk
+        logits = logits_list[idx][:, start_cut:end_cut]
         cutted_logits.append(logits)
 
     # join