Modified app.py, added streamlit session state persistence

app.py

@@ -5,21 +5,15 @@ import streamlit as st
 import torchaudio as ta
 
 from io import BytesIO
-from transformers import AutoProcessor, SeamlessM4TModel, WhisperProcessor, WhisperForConditionalGeneration
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
-if torch.cuda.is_available():
-    device = "cuda:0"
-    torch_dtype = torch.float16
-else:
-    device = "cpu"
-    torch_dtype = torch.float32
+# Set up device and dtype
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if device == "cuda:0" else torch.float32
 
-SAMPLING_RATE=16000
-task = "transcribe"
+SAMPLING_RATE = 16000
 
-print(f"{device} Active!")
-
-# load Whisper model and processor
+# Load Whisper model and processor
 processor = WhisperProcessor.from_pretrained("openai/whisper-small")
 model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
 
@@ -31,113 +25,74 @@ st.sidebar.header("Upload Audio Files")
 uploaded_files = st.sidebar.file_uploader("Choose audio files", type=["mp3", "wav"], accept_multiple_files=True)
 submit_button = st.sidebar.button("Submit")
 
+# Session state to hold data
+if 'audio_files' not in st.session_state:
+    st.session_state.audio_files = []
+    st.session_state.transcriptions = {}
+    st.session_state.translations = {}
+    st.session_state.detected_languages = []
+    st.session_state.waveforms = []
 
-# def transcribe_audio(audio_data):
-#     recognizer = sr.Recognizer()
-#     with sr.AudioFile(audio_data) as source:
-#         audio = recognizer.record(source)
-#     try:
-#         # Transcribe the audio using Google Web Speech API
-#         transcription = recognizer.recognize_google(audio)
-#         return transcription
-#     except sr.UnknownValueError:
-#         return "Unable to transcribe the audio."
-#     except sr.RequestError as e:
-#         return f"Could not request results; {e}"
 
 def detect_language(audio_file):
-    whisper_model = whisper.load_model("base")
+    whisper_model = whisper.load_model("small")
+    trimmed_audio = whisper.pad_or_trim(audio_file)
     mel = whisper.log_mel_spectrogram(trimmed_audio).to(whisper_model.device)
-    # detect the spoken language
-    _, probs = whisper_model.detect_language(mel)
-    print(f"Detected language: {max(probs[0], key=probs[0].get)}")
-    return max(probs[0], key=probs[0].get)
-    
-# if submit_button and uploaded_files is not None:
-#     st.write("Files uploaded successfully!")
-
-#     for uploaded_file in uploaded_files:
-#         # Display file name and audio player
-
-#         st.write(f"**File name**: {uploaded_file.name}")
-#         st.audio(uploaded_file, format=uploaded_file.type)
-
-#         # Transcription section
-#         st.write("**Transcription**:")
-
-#         # Read the uploaded file data
-#         waveform, sampling_rate = ta.load(uploaded_file.getvalue())
-#         resampled_inp = ta.functional.resample(waveform, orig_freq=sampling_rate, new_freq=SAMPLING_RATE)
-
-#         input_features = processor(resampled_inp[0], sampling_rate=16000, return_tensors='pt').input_features
-
-#         if task == "translate":
-            
-#             # Detect Language 
-#             lang = detect_language(input_features)
-#             with open('languages.pkl', 'rb') as f:
-#                 lang_dict = pickle.load(f)
-#             detected_language = lang_dict[lang]
-
-#             # Set decoder & Predict translation
-#             forced_decoder_ids = processor.get_decoder_prompt_ids(language=detected_language, task="translate")
-#             predicted_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
-#         else:
-#             predicted_ids = model.generate(input_features)
-#         # decode token ids to text
-#         transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
-#         for i in range(len(transcription)):
-#             st.write(transcription[i])
-#         # print(waveform, sampling_rate)
-#         # Run transcription function and display
-#         # import pdb;pdb.set_trace()
-#         # st.write(audio_data.getvalue())
-
+    _, probs = whisper_model.detect_language(mel[0])
+    detected_lang = max(probs, key=probs.get)
+    print(f"Detected language: {detected_lang}")
+    return detected_lang
 
 
+# Process uploaded files
 if submit_button and uploaded_files is not None:
-    # Initialize a list to store detected languages
-    detected_languages = []
+    st.session_state.audio_files = uploaded_files
+    st.session_state.detected_languages = []
 
     for uploaded_file in uploaded_files:
-        # Read the uploaded file data
         waveform, sampling_rate = ta.load(BytesIO(uploaded_file.read()))
-
-        # Resample if necessary
         if sampling_rate != SAMPLING_RATE:
             waveform = ta.functional.resample(waveform, orig_freq=sampling_rate, new_freq=SAMPLING_RATE)
 
-        # Detect language
-        detected_language = detect_language(waveform, SAMPLING_RATE)
-        detected_languages.append(detected_language)
+        st.session_state.waveforms.append(waveform)
+        detected_language = detect_language(waveform)
+        st.session_state.detected_languages.append(detected_language)
 
-    # Display each uploaded file with its detected language and an audio player
-    for i, uploaded_file in enumerate(uploaded_files):
-        col1, col2 = st.columns([1, 3])  # Two columns, one for the player, one for the buttons
+# Display uploaded files and options
+if 'audio_files' in st.session_state and st.session_state.audio_files:
+    for i, uploaded_file in enumerate(st.session_state.audio_files):
+        col1, col2 = st.columns([1, 3])
 
         with col1:
             st.write(f"**File name**: {uploaded_file.name}")
-            st.audio(BytesIO(uploaded_file.getvalue()), format=uploaded_file.type)
-            st.write(f"**Detected Language**: {detected_languages[i]}")
+            st.audio(BytesIO(uploaded_file.read()), format=uploaded_file.type)
+            st.write(f"**Detected Language**: {st.session_state.detected_languages[i]}")
 
         with col2:
-            # Add Transcription and Translation buttons
+            # 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}"):
-                # Transcription process
-                input_features = processor(waveform[0], sampling_rate=SAMPLING_RATE, return_tensors='pt').input_features
                 predicted_ids = model.generate(input_features)
                 transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)
-                for line in transcription:
+                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)
 
             if st.button(f"Translate {uploaded_file.name}"):
-                # Translation process
                 with open('languages.pkl', 'rb') as f:
                     lang_dict = pickle.load(f)
-                detected_language_name = lang_dict[detected_languages[i]]
+                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)
-                for line in translation:
+                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)