Implemented Chunking

app.py

@@ -3,6 +3,7 @@ import pickle
 import whisper
 import streamlit as st
 import torchaudio as ta
+import numpy as np
 
 from io import BytesIO
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
@@ -12,10 +13,11 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda:0" else torch.float32
 
 SAMPLING_RATE = 16000
+CHUNK_LENGTH_S = 20  # 30 seconds per chunk
 
 # Load Whisper model and processor
 processor = WhisperProcessor.from_pretrained("openai/whisper-small")
-model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
+model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
 
 # Title of the app
 st.title("Audio Player with Live Transcription")
@@ -36,18 +38,42 @@ if 'audio_files' not in st.session_state:
 
 def detect_language(audio_file):
     whisper_model = whisper.load_model("small")
-    trimmed_audio = whisper.pad_or_trim(audio_file)
+    trimmed_audio = whisper.pad_or_trim(audio_file.squeeze())
     mel = whisper.log_mel_spectrogram(trimmed_audio).to(whisper_model.device)
-    _, probs = whisper_model.detect_language(mel[0])
-    detected_lang = max(probs, key=probs.get)
+    _, probs = whisper_model.detect_language(mel)
+    detected_lang = max(probs[0], key=probs[0].get)
     print(f"Detected language: {detected_lang}")
     return detected_lang
 
 
+def process_long_audio(waveform, sampling_rate, task="transcribe", language=None):
+    input_length = waveform.shape[1]
+    chunk_length = int(CHUNK_LENGTH_S * sampling_rate)
+    chunks = [waveform[:, i:i + chunk_length] for i in range(0, input_length, chunk_length)]
+
+    results = []
+    for chunk in chunks:
+        # import pdb;pdb.set_trace()
+        input_features = processor(chunk[0], sampling_rate=sampling_rate, return_tensors="pt").input_features.to(device)
+
+        with torch.no_grad():
+            if task == "translate":
+                forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task="translate")
+                generated_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
+            else:
+                generated_ids = model.generate(input_features)
+
+        transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)
+        results.extend(transcription)
+
+    return " ".join(results)
+
+
 # Process uploaded files
 if submit_button and uploaded_files is not None:
     st.session_state.audio_files = uploaded_files
     st.session_state.detected_languages = []
+    st.session_state.waveforms = []
 
     for uploaded_file in uploaded_files:
         waveform, sampling_rate = ta.load(BytesIO(uploaded_file.read()))
@@ -69,30 +95,25 @@ if 'audio_files' in st.session_state and st.session_state.audio_files:
             st.write(f"**Detected Language**: {st.session_state.detected_languages[i]}")
 
         with col2:
-            # import pdb;pdb.set_trace()
-            input_features = processor(st.session_state.waveforms[i][0], sampling_rate=SAMPLING_RATE, return_tensors='pt').input_features
-
             if st.button(f"Transcribe {uploaded_file.name}"):
-                predicted_ids = model.generate(input_features)
-                transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
-                st.session_state.transcriptions[i] = transcription
+                with st.spinner("Transcribing..."):
+                    transcription = process_long_audio(st.session_state.waveforms[i], SAMPLING_RATE)
+                    st.session_state.transcriptions[i] = transcription
 
             if st.session_state.transcriptions.get(i):
                 st.write("**Transcription**:")
-                for line in st.session_state.transcriptions[i]:
-                    st.write(line)
+                st.write(st.session_state.transcriptions[i])
 
             if st.button(f"Translate {uploaded_file.name}"):
-                with open('languages.pkl', 'rb') as f:
-                    lang_dict = pickle.load(f)
-                detected_language_name = lang_dict[st.session_state.detected_languages[i]]
+                with st.spinner("Translating..."):
+                    with open('languages.pkl', 'rb') as f:
+                        lang_dict = pickle.load(f)
+                    detected_language_name = lang_dict[st.session_state.detected_languages[i]]
 
-                forced_decoder_ids = processor.get_decoder_prompt_ids(language=detected_language_name, task="translate")
-                predicted_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
-                translation = processor.batch_decode(predicted_ids, skip_special_tokens=True)
-                st.session_state.translations[i] = translation
+                    translation = process_long_audio(st.session_state.waveforms[i], SAMPLING_RATE, task="translate",
+                                                     language=detected_language_name)
+                    st.session_state.translations[i] = translation
 
             if st.session_state.translations.get(i):
                 st.write("**Translation**:")
-                for line in st.session_state.translations[i]:
-                    st.write(line)
+                st.write(st.session_state.translations[i])