Update custom_interface_app.py

custom_interface_app.py

@@ -41,7 +41,7 @@ class ASR(Pretrained):
         # Forward encoder + decoder
         tokens = torch.tensor([[1, 1]]) * self.mods.whisper.config.decoder_start_token_id
         tokens = tokens.to(device)
-        enc_out, logits, _ = self.mods.whisper(wavs.detach(), tokens.detach())
+        enc_out, logits, _ = self.mods.whisper(wavs, tokens)
         log_probs = self.hparams.log_softmax(logits)
 
         hyps, _, _, _ = self.hparams.test_search(enc_out.detach(), wav_lens)
@@ -128,30 +128,53 @@ class ASR(Pretrained):
 
     def classify_file_w2v2(self, waveform, device):
         # Load the audio file
+        # path = "long_sample.wav"
         # waveform, sr = librosa.load(path, sr=16000)
 
+        # increase the volume if needed
+        # waveform = self.increase_volume(waveform)
+
         # Get audio length in seconds
-        audio_length = len(waveform) / 16000
+        sr = 16000
+        audio_length = len(waveform) / sr
         
-        if audio_length >= 30:
-            # split audio every 20 seconds
+        if audio_length >= 20:
+            print(f"Audio is too long ({audio_length:.2f} seconds), splitting into segments")
+            # Detect non-silent segments
+
+            non_silent_intervals = librosa.effects.split(waveform, top_db=20)  # Adjust top_db for sensitivity
+
             segments = []
-            all_segments = []
-            max_duration = 30 * 16000  # Maximum segment duration in samples (20 seconds)
-            num_segments = int(np.ceil(len(waveform) / max_duration))
-            start = 0
-            for i in range(num_segments):
-                end = start + max_duration
-                if end > len(waveform):
-                    end = len(waveform)
+            current_segment = []
+            current_length = 0
+            max_duration = 20 * sr  # Maximum segment duration in samples (20 seconds)
+
+
+            for interval in non_silent_intervals:
+                start, end = interval
                 segment_part = waveform[start:end]
-                segment_len = len(segment_part) / 16000
-                if segment_len < 1:
-                    continue
-                segments.append(segment_part)
-                start = end
 
-            for segment in segments:
+                # If adding the next part exceeds max duration, store the segment and start a new one
+                if current_length + len(segment_part) > max_duration:
+                    segments.append(np.concatenate(current_segment))
+                    current_segment = []
+                    current_length = 0
+
+                current_segment.append(segment_part)
+                current_length += len(segment_part)
+
+            # Append the last segment if it's not empty
+            if current_segment:
+                segments.append(np.concatenate(current_segment))
+
+            # Process each segment
+            outputs = []
+            for i, segment in enumerate(segments):
+                print(f"Processing segment {i + 1}/{len(segments)}, length: {len(segment) / sr:.2f} seconds")
+
+                # import soundfile as sf
+                # sf.write(f"outputs/segment_{i}.wav", segment, sr)
+
                 segment_tensor = torch.tensor(segment).to(device)
 
                 # Fake a batch for the segment
@@ -159,171 +182,145 @@ class ASR(Pretrained):
                 rel_length = torch.tensor([1.0]).to(device)  # Adjust if necessary
 
                 # Pass the segment through the ASR model
-                segment_output = self.encode_batch_w2v2(device, batch, rel_length)
-                segment_output = [" ".join(segment) for segment in segment_output]
-                all_segments.append(segment_output)
-            
-            segments = ""
-            for segment in all_segments:
-                segment = segment[0]
-                segments += segment + " "
-            return [segments]
+                result = " ".join(self.encode_batch_w2v2(device, batch, rel_length)[0])
+                outputs.append(result)
+            return outputs
         else:
             waveform = torch.tensor(waveform).to(device)
             waveform = waveform.to(device)
             # Fake a batch:
             batch = waveform.unsqueeze(0)
             rel_length = torch.tensor([1.0]).to(device)
-            outputs = self.encode_batch_w2v2(device, batch, rel_length)
-            return [" ".join(out) for out in outputs]
-    
+            outputs = " ".join(self.encode_batch_w2v2(device, batch, rel_length)[0])
+            return [outputs]
+
 
     def classify_file_whisper_mkd(self, waveform, device):
         # Load the audio file
+        # path = "long_sample.wav"
         # waveform, sr = librosa.load(path, sr=16000)
 
+        # increase the volume if needed
+        # waveform = self.increase_volume(waveform)
+
         # Get audio length in seconds
-        audio_length = len(waveform) / 16000
+        sr = 16000
+        audio_length = len(waveform) / sr
         
-        if audio_length >= 30:
-            # split audio every 20 seconds
+        if audio_length >= 20:
+            print(f"Audio is too long ({audio_length:.2f} seconds), splitting into segments")
+            # Detect non-silent segments
+
+            non_silent_intervals = librosa.effects.split(waveform, top_db=20)  # Adjust top_db for sensitivity
+
             segments = []
-            all_segments = []
-            max_duration = 30 * 16000  # Maximum segment duration in samples (20 seconds)
-            num_segments = int(np.ceil(len(waveform) / max_duration))
-            start = 0
-            for i in range(num_segments):
-                end = start + max_duration
-                if end > len(waveform):
-                    end = len(waveform)
-                segment_part = waveform[start:end]
-                segment_len = len(segment_part) / 16000
-                if segment_len < 1:
-                    continue
-                segments.append(segment_part)
-                start = end
+            current_segment = []
+            current_length = 0
+            max_duration = 20 * sr  # Maximum segment duration in samples (20 seconds)
 
-            for segment in segments:
-                segment_tensor = torch.tensor(segment).to(device)
+            for interval in non_silent_intervals:
+                start, end = interval
+                segment_part = waveform[start:end]
 
-                # Fake a batch for the segment
-                batch = segment_tensor.unsqueeze(0).to(device)
-                rel_length = torch.tensor([1.0]).to(device)
+                # If adding the next part exceeds max duration, store the segment and start a new one
+                if current_length + len(segment_part) > max_duration:
+                    segments.append(np.concatenate(current_segment))
+                    current_segment = []
+                    current_length = 0
 
-                # Pass the segment through the ASR model
-                segment_output = self.encode_batch_whisper(device, batch, rel_length)
-                # segment_output = [" ".join(segment) for segment in segment_output]
-                all_segments.append(segment_output)
-            
-            segments = ""
-            for segment in all_segments:
-                segment = segment[0]
-                segments += segment + " "
-            return [segments]
-        else:
-            waveform = torch.tensor(waveform).to(device)
-            waveform = waveform.to(device)
-            batch = waveform.unsqueeze(0)
-            rel_length = torch.tensor([1.0]).to(device)
-            outputs = self.encode_batch_whisper(device, batch, rel_length)
-            return outputs
+                current_segment.append(segment_part)
+                current_length += len(segment_part)
 
+            # Append the last segment if it's not empty
+            if current_segment:
+                segments.append(np.concatenate(current_segment))
 
-    def classify_file_whisper_mkd_streaming(self, waveform, device):
-        # Load the audio file
-        # waveform, sr = librosa.load(path, sr=16000)
+            # Process each segment
+            outputs = []
+            for i, segment in enumerate(segments):
+                print(f"Processing segment {i + 1}/{len(segments)}, length: {len(segment) / sr:.2f} seconds")
 
-        # Get audio length in seconds
-        audio_length = len(waveform) / 16000
-        
-        if audio_length >= 30:
-            # split audio every 30 seconds
-            segments = []
-            max_duration = 30 * 16000  # Maximum segment duration in samples (20 seconds)
-            num_segments = int(np.ceil(len(waveform) / max_duration))
-            start = 0
-            for i in range(num_segments):
-                end = start + max_duration
-                if end > len(waveform):
-                    end = len(waveform)
-                segment_part = waveform[start:end]
-                segment_len = len(segment_part) / 16000
-                if segment_len < 1:
-                    continue
-                segments.append(segment_part)
-                start = end
+                # import soundfile as sf
+                # sf.write(f"outputs/segment_{i}.wav", segment, sr)
 
-            for segment in segments:
                 segment_tensor = torch.tensor(segment).to(device)
 
                 # Fake a batch for the segment
                 batch = segment_tensor.unsqueeze(0).to(device)
-                rel_length = torch.tensor([1.0]).to(device)
+                rel_length = torch.tensor([1.0]).to(device)  # Adjust if necessary
 
                 # Pass the segment through the ASR model
                 segment_output = self.encode_batch_whisper(device, batch, rel_length)
-                yield segment_output
+                outputs.append(segment_output)
+            return outputs
         else:
             waveform = torch.tensor(waveform).to(device)
             waveform = waveform.to(device)
+            # Fake a batch:
             batch = waveform.unsqueeze(0)
             rel_length = torch.tensor([1.0]).to(device)
-            outputs = self.encode_batch_whisper(device, batch, rel_length)
-            yield outputs
+            outputs.append(self.encode_batch_whisper(device, batch, rel_length))
+            return outputs
 
 
-    def classify_file_whisper(self, waveform, pipe, device):
-        # waveform, sr = librosa.load(path, sr=16000)
+    def classify_file_whisper(self, path, pipe, device):
+        waveform, sr = librosa.load(path, sr=16000)
         transcription = pipe(waveform, generate_kwargs={"language": "macedonian"})["text"]
-        return [transcription]
+        return transcription
        
 
-    def classify_file_mms(self, waveform, processor, model, device):
+    def classify_file_mms(self, path, processor, model, device):
         # Load the audio file
-        # waveform, sr = librosa.load(path, sr=16000)
+        waveform, sr = librosa.load(path, sr=16000)
 
         # Get audio length in seconds
-        audio_length = len(waveform) / 16000
+        audio_length = len(waveform) / sr
         
-        if audio_length >= 30:
-            # split audio every 20 seconds
+        if audio_length >= 20:
+            print(f"MMS Audio is too long ({audio_length:.2f} seconds), splitting into segments")
+            # Detect non-silent segments
+            non_silent_intervals = librosa.effects.split(waveform, top_db=20)  # Adjust top_db for sensitivity
+
             segments = []
-            all_segments = []
-            max_duration = 30 * 16000  # Maximum segment duration in samples (20 seconds)
-            num_segments = int(np.ceil(len(waveform) / max_duration))
-            start = 0
-            for i in range(num_segments):
-                end = start + max_duration
-                if end > len(waveform):
-                    end = len(waveform)
+            current_segment = []
+            current_length = 0
+            max_duration = 20 * sr  # Maximum segment duration in samples (20 seconds)
+
+
+            for interval in non_silent_intervals:
+                start, end = interval
                 segment_part = waveform[start:end]
-                segment_len = len(segment_part) / 16000
-                if segment_len < 1:
-                    continue
-                segments.append(segment_part)
-                start = end
 
-            for segment in segments:
+                # If adding the next part exceeds max duration, store the segment and start a new one
+                if current_length + len(segment_part) > max_duration:
+                    segments.append(np.concatenate(current_segment))
+                    current_segment = []
+                    current_length = 0
+
+                current_segment.append(segment_part)
+                current_length += len(segment_part)
+
+            # Append the last segment if it's not empty
+            if current_segment:
+                segments.append(np.concatenate(current_segment))
+
+            # Process each segment
+            outputs = []
+            for i, segment in enumerate(segments):
+                print(f"MMS Processing segment {i + 1}/{len(segments)}, length: {len(segment) / sr:.2f} seconds")
+
                 segment_tensor = torch.tensor(segment).to(device)
 
                 # Pass the segment through the ASR model
                 inputs = processor(segment_tensor, sampling_rate=16_000, return_tensors="pt").to(device)
-                inputs['input_values'] = inputs['input_values']
                 outputs = model(**inputs).logits
                 ids = torch.argmax(outputs, dim=-1)[0]
                 segment_output = processor.decode(ids)
-                # segment_output = [" ".join(segment) for segment in segment_output]
-                all_segments.append(segment_output)
-            
-            segments = ""
-            for segment in all_segments:
-                segments += segment + " "
-            return [segments]
+                yield segment_output
         else:
             waveform = torch.tensor(waveform).to(device)
             inputs = processor(waveform, sampling_rate=16_000, return_tensors="pt").to(device)
-            inputs['input_values'] = inputs['input_values']
             outputs = model(**inputs).logits
             ids = torch.argmax(outputs, dim=-1)[0]
             transcription = processor.decode(ids)
-            return [transcription]
+            yield transcription