Added decoding at the end

app.py

@@ -29,24 +29,30 @@ def resample(sr, audio_data):
     return audio_16k
 
 
-def model(audio_16k, is_start):
+def model(audio_16k):
     logits, logits_len, greedy_predictions = asr_model.forward(
         input_signal=torch.tensor([audio_16k]), 
         input_signal_length=torch.tensor([len(audio_16k)])
     )
-
-    # cut overhead
-    buffer_len = len(audio_16k)
-    logits_overhead = (logits.shape[1] - 1) * overhead_len // buffer_len
-    logits_overhead //= 2
-    delay = (logits.shape[1] - 1) - (2 * logits_overhead)
-    start_cut = 0 if is_start else logits_overhead
-    delay += 0 if not is_start else logits_overhead
-    logits = logits[:, start_cut:start_cut+delay]
     return logits
 
 
-def decode_predictions(logits):
+def decode_predictions(logits_list):
+    # calc overhead
+    logits_overhead = logits_list[0].shape[1] * overhead_len // total_buffer
+    logits_overhead //= 2
+    #delay = (logits.shape[1] - 1) - (2 * logits_overhead)
+
+    # cut overhead
+    cutted_logits = []
+    for idx in range(len(logits_list)):
+        start_cut = 0 if (idx==0) else logits_overhead
+        end_cut = 1 if (idx==len(logits_list)-1) else logits_overhead
+        logits = logits_list[idx][:, start_cut:-end_cut]
+        cutted_logits.append(logits)
+
+    # join
+    logits = torch.cat(cutted_logits, axis=1)
     logits_len = torch.tensor([logits.shape[1]])
     current_hypotheses, all_hyp = asr_model.decoding.ctc_decoder_predictions_tensor(
         logits, decoder_lengths=logits_len, return_hypotheses=False,
@@ -57,8 +63,7 @@ def decode_predictions(logits):
 
 def transcribe(audio, state):
     if state is None:
-        state = [np.array([], dtype=np.float32), None]
-    is_start = state[1] is None
+        state = [np.array([], dtype=np.float32), []]
 
     sr, audio_data = audio
     audio_16k = resample(sr, audio_data)
@@ -70,15 +75,11 @@ def transcribe(audio, state):
         buffer = state[0][:total_buffer]
         state[0] = state[0][total_buffer - overhead_len:]
         # run model
-        is_start = state[1] is None
-        logits = model(buffer, is_start)
+        logits = model(buffer)
         # add logits
-        if is_start:
-            state[1] = logits
-        else:
-            state[1] = torch.cat([state[1],logits], axis=1)
+        state[1].append(logits)
 
-    if is_start:
+    if len(state[1]) == 0:
         text = ""
     else:
         text = decode_predictions(state[1])