Update custom_interface_app.py

custom_interface_app.py

@@ -8,7 +8,7 @@ class ASR(Pretrained):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def encode_batch(self, device, wavs, wav_lens=None, normalize=False):
+    def encode_batch_w2v2(self, device, wavs, wav_lens=None, normalize=False):
         wavs = wavs.to(device)
         wav_lens = wav_lens.to(device)
 
@@ -33,6 +33,22 @@ class ASR(Pretrained):
         predicted_words = prediction
         return predicted_words
 
+
+    def encode_batch_whisper(self, device, wavs, wav_lens=None, normalize=False):
+        wavs = wavs.to(device)
+        wav_lens = wav_lens.to(device)
+
+        # 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, tokens)
+        log_probs = self.hparams.log_softmax(logits)
+
+        hyps, _, _, _ = self.hparams.test_search(enc_out.detach(), wav_lens)
+        predicted_words = [self.mods.whisper.tokenizer.decode(token, skip_special_tokens=True).strip() for token in hyps]
+        return predicted_words
+
+
     def filter_repetitions(self, seq, max_repetition_length):
         seq = list(seq)
         output = []
@@ -110,7 +126,74 @@ class ASR(Pretrained):
         return waveform
 
 
-    def classify_file(self, path, device):
+    def classify_file_w2v2(self, path, 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) / sr
+        
+        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 = []
+            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]
+
+                # 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
+                batch = segment_tensor.unsqueeze(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_w2v2(device, batch, rel_length)
+                yield segment_output
+        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)
+            yield outputs
+
+
+    def classify_file_whisper_mkd(self, path, device):
         # Load the audio file
         # path = "long_sample.wav"
         waveform, sr = librosa.load(path, sr=16000)
@@ -165,7 +248,7 @@ 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(device, batch, rel_length)
+                segment_output = self.encode_batch_whisper(device, batch, rel_length)
                 yield segment_output
         else:
             waveform = torch.tensor(waveform).to(device)
@@ -173,7 +256,7 @@ class ASR(Pretrained):
             # Fake a batch:
             batch = waveform.unsqueeze(0)
             rel_length = torch.tensor([1.0]).to(device)
-            outputs = self.encode_batch(device, batch, rel_length)
+            outputs = self.encode_batch_whisper(device, batch, rel_length)
             yield outputs