Update app.py

app.py

@@ -1,27 +1,30 @@
 import streamlit as st
 import os
+import librosa
+import torch
 from pydub import AudioSegment
 from pydub.silence import split_on_silence
 from dotenv import load_dotenv
 from tempfile import NamedTemporaryFile
 import math
 from docx import Document
-import whisper
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
-# Load environment variables from .env file (if needed for other configurations)
+# Load environment variables from .env file (if needed)
 load_dotenv()
 
 @st.cache_resource
 def load_whisper_model():
     """
-    Load the Whisper model once and cache it for future use.
-    You can choose the model size: "tiny", "base", "small", "medium", or "large".
+    Load the Whisper model and processor from Hugging Face.
+    You can change the model variant ("openai/whisper-base" is used here).
     """
-    model = whisper.load_model("base")
-    return model
+    model_name = "openai/whisper-base"  # Options: "tiny", "base", "small", "medium", "large"
+    processor = WhisperProcessor.from_pretrained(model_name)
+    model = WhisperForConditionalGeneration.from_pretrained(model_name)
+    return processor, model
 
-# Load the Whisper model globally so it’s only loaded once.
-model = load_whisper_model()
+processor, model = load_whisper_model()
 
 def split_audio_on_silence(audio_file_path, min_silence_len=500, silence_thresh=-40, keep_silence=250):
     """
@@ -47,7 +50,7 @@ def split_audio_on_silence(audio_file_path, min_silence_len=500, silence_thresh=
 
 def transcribe(audio_file):
     """
-    Transcribe an audio file using the locally loaded Whisper model.
+    Transcribe an audio file using the locally loaded Whisper model from Hugging Face.
     
     Args:
         audio_file (str): Path to the audio file.
@@ -55,12 +58,17 @@ def transcribe(audio_file):
     Returns:
         str: Transcribed text.
     """
-    result = model.transcribe(audio_file, language="en")
-    return result["text"]
+    # Load audio using librosa, resampling to 16000 Hz as required by Whisper
+    speech, sr = librosa.load(audio_file, sr=16000)
+    input_features = processor(speech, sampling_rate=16000, return_tensors="pt").input_features
+    # Generate transcription
+    predicted_ids = model.generate(input_features)
+    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+    return transcription
 
 def process_audio_chunks(audio_chunks):
     """
-    Process and transcribe each audio chunk in sequence.
+    Process and transcribe each audio chunk.
     
     Args:
         audio_chunks (list): List of AudioSegment chunks.
@@ -69,22 +77,22 @@ def process_audio_chunks(audio_chunks):
         str: Combined transcription from all chunks.
     """
     transcriptions = []
-    min_length_ms = 100  # Minimum length required for processing
+    min_length_ms = 100  # Minimum length required (0.1 seconds)
+    
     for i, chunk in enumerate(audio_chunks):
         if len(chunk) < min_length_ms:
             st.warning(f"Chunk {i} is too short to be processed.")
             continue
         
-        # Save the chunk temporarily as a WAV file
         with NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio_file:
             chunk.export(temp_audio_file.name, format="wav")
             temp_audio_file_path = temp_audio_file.name
-
+        
         transcription = transcribe(temp_audio_file_path)
         if transcription:
             transcriptions.append(transcription)
             st.write(f"Transcription for chunk {i}: {transcription}")
-
+        
         os.remove(temp_audio_file_path)
     return " ".join(transcriptions)
 
@@ -106,7 +114,7 @@ def save_transcription_to_docx(transcription, audio_file_path):
     doc.save(output_file_name)
     return output_file_name
 
-st.title("Audio Transcription with Whisper (Local)")
+st.title("Audio Transcription with Whisper (Local via Hugging Face)")
 
 # Allow uploading of audio or video files
 uploaded_file = st.file_uploader("Upload an audio or video file", type=["wav", "mp3", "ogg", "m4a", "mp4", "mov"])
@@ -122,14 +130,14 @@ if uploaded_file is not None and st.session_state.transcription is None:
     temp_audio_file = f"temp_audio_file.{file_extension}"
     with open(temp_audio_file, "wb") as f:
         f.write(uploaded_file.getbuffer())
-
+    
     # Split and process audio using silence detection
     with st.spinner('Transcribing...'):
         audio_chunks = split_audio_on_silence(
             temp_audio_file, 
-            min_silence_len=500,   # adjust based on your audio
-            silence_thresh=-40,    # adjust based on ambient noise level
-            keep_silence=250       # retains a bit of silence at the edges
+            min_silence_len=500,
+            silence_thresh=-40,
+            keep_silence=250
         )
         transcription = process_audio_chunks(audio_chunks)
         if transcription:
@@ -137,7 +145,7 @@ if uploaded_file is not None and st.session_state.transcription is None:
             st.success('Transcription complete!')
             output_docx_file = save_transcription_to_docx(transcription, uploaded_file.name)
             st.session_state.output_docx_file = output_docx_file
-            
+    
     if os.path.exists(temp_audio_file):
         os.remove(temp_audio_file)