Nearly fixed streaming bug

app.py

@@ -17,7 +17,7 @@ asr_model.decoder.freeze()
 
 
 total_buffer = asr_model.cfg["sample_rate"]
-overhead_len = asr_model.cfg["sample_rate"] // 4
+overhead_len = asr_model.cfg["sample_rate"] // 2
 model_stride = 4
 
 
@@ -29,16 +29,20 @@ def resample(sr, audio_data):
     return audio_16k
 
 
-def model(audio_16k):
+def model(audio_16k, is_start):
     logits, logits_len, greedy_predictions = asr_model.forward(
         input_signal=torch.tensor([audio_16k]), 
         input_signal_length=torch.tensor([len(audio_16k)])
     )
 
     # cut overhead
-    logits_overhead = logits.shape[1] * overhead_len // total_buffer
-    extra = 1 if (logits.shape[1] * overhead_len % total_buffer) else 0
-    logits = logits[:,logits_overhead:-logits_overhead-extra]
+    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
 
 
@@ -54,6 +58,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
 
     sr, audio_data = audio
     audio_16k = resample(sr, audio_data)
@@ -61,20 +66,22 @@ def transcribe(audio, state):
     # join to audio sequence
     state[0] = np.concatenate([state[0], audio_16k])
 
-    buffer_len = len(state[0])
-    if (buffer_len > total_buffer):
-        buffer_len = buffer_len - buffer_len % total_buffer
-        buffer = state[0][:buffer_len]
-        state[0] = state[0][buffer_len - overhead_len:]
+    while (len(state[0]) > total_buffer):
+        buffer = state[0][:total_buffer]
+        state[0] = state[0][total_buffer - overhead_len:]
         # run model
-        logits = model(buffer)
+        is_start = state[1] is None
+        logits = model(buffer, is_start)
         # add logits
-        if state[1] is None:
+        if is_start:
             state[1] = logits
         else:
             state[1] = torch.cat([state[1],logits], axis=1)
 
-    text = decode_predictions(state[1])
+    if is_start:
+        text = ""
+    else:
+        text = decode_predictions(state[1])
     return text, state