Update app.py

app.py

@@ -7,126 +7,163 @@ Date: January 2025
 """
 
 import streamlit as st
-from pyannote.audio import Pipeline
-import whisper
-import tempfile
+from src.models.diarization import SpeakerDiarizer
+from src.models.transcription import Transcriber
+from src.models.summarization import Summarizer
+from src.utils.audio_processor import AudioProcessor
+from src.utils.formatter import TimeFormatter
 import os
-import torch
-from transformers import pipeline as tf_pipeline, BartTokenizer
-from pydub import AudioSegment
-import io
-import pickle
-
-class SpeakerDiarizer:
-   def __init__(self, token):
-       self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=token)
-
-   def process(self, audio_file):
-       return self.pipeline(audio_file)
-
-class Transcriber:
-   def __init__(self):
-       self.model = whisper.load_model("base")
-
-   def process(self, audio_file):
-       return self.model.transcribe(audio_file)["text"]
-
-class Summarizer:
-   def __init__(self, model_path='bart_ami_finetuned.pkl'):
-       self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
-       with open(model_path, 'rb') as f:
-           self.model = pickle.load(f)
-
-   def process(self, text):
-       inputs = self.tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
-       summary_ids = self.model.generate(inputs["input_ids"], max_length=150, min_length=40)
-       return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
+# Cache for model loading
 @st.cache_resource
 def load_models():
-   try:
-       diarizer = SpeakerDiarizer(st.secrets["hf_token"])
-       transcriber = Transcriber()
-       summarizer = Summarizer()
-       return diarizer, transcriber, summarizer
-   except Exception as e:
-       st.error(f"Error loading models: {str(e)}")
-       return None, None, None
-
-def process_audio(audio_file):
-   try:
-       audio_bytes = io.BytesIO(audio_file.getvalue())
-       
-       with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-           if audio_file.name.lower().endswith('.mp3'):
-               audio = AudioSegment.from_mp3(audio_bytes)
-           else:
-               audio = AudioSegment.from_wav(audio_bytes)
-           
-           audio = audio.set_frame_rate(16000).set_channels(1).set_sample_width(2)
-           audio.export(tmp.name, format="wav", parameters=["-ac", "1", "-ar", "16000"])
-           tmp_path = tmp.name
-
-           diarizer, transcriber, summarizer = load_models()
-           if not all([diarizer, transcriber, summarizer]):
-               return "Model loading failed"
-
-           with st.spinner("Processing..."):
-               diarization = diarizer.process(tmp_path)
-               transcription = transcriber.process(tmp_path)
-               summary = summarizer.process(transcription)
-
-           os.unlink(tmp_path)
-           
-           return {
-               "diarization": diarization,
-               "transcription": transcription,
-               "summary": summary
-           }
-           
-   except Exception as e:
-       st.error(f"Error: {str(e)}")
-       return None
+    """
+    Load and cache all required models.
+    
+    Returns:
+        tuple: (diarizer, transcriber, summarizer) or (None, None, None) if loading fails
+    """
+    try:
+        diarizer = SpeakerDiarizer(st.secrets["hf_token"])
+        diarizer_model = diarizer.load_model()
+        
+        transcriber = Transcriber()
+        transcriber_model = transcriber.load_model()
+        
+        summarizer = Summarizer()
+        summarizer_model = summarizer.load_model()
+        
+        if not all([diarizer_model, transcriber_model, summarizer_model]):
+            raise ValueError("One or more models failed to load")
+            
+        return diarizer, transcriber, summarizer
+    except Exception as e:
+        st.error(f"Error loading models: {str(e)}")
+        st.error("Debug info: Check if HF token is valid and has necessary permissions")
+        return None, None, None
+
+def process_audio(audio_file, max_duration=600):
+    """
+    Process the uploaded audio file through all models.
+    
+    Args:
+        audio_file: Uploaded audio file
+        max_duration (int): Maximum duration in seconds
+        
+    Returns:
+        dict: Processing results containing diarization, transcription, and summary
+    """
+    try:
+        # Process audio file
+        audio_processor = AudioProcessor()
+        tmp_path = audio_processor.standardize_audio(audio_file)
+        
+        # Load models
+        diarizer, transcriber, summarizer = load_models()
+        if not all([diarizer, transcriber, summarizer]):
+            return "Model loading failed"
+
+        # Process with each model
+        with st.spinner("Identifying speakers..."):
+            diarization_result = diarizer.process(tmp_path)
+        
+        with st.spinner("Transcribing audio..."):
+            transcription = transcriber.process(tmp_path)
+            
+        with st.spinner("Generating summary..."):
+            summary = summarizer.process(transcription["text"])
+
+        # Cleanup
+        os.unlink(tmp_path)
+        
+        return {
+            "diarization": diarization_result,
+            "transcription": transcription,
+            "summary": summary[0]["summary_text"]
+        }
+        
+    except Exception as e:
+        st.error(f"Error processing audio: {str(e)}")
+        return None
 
 def main():
-   st.title("Multi-Speaker Audio Analyzer")
-   st.write("Upload an audio file (MP3/WAV) up to 5 minutes long")
-
-   uploaded_file = st.file_uploader("Choose a file", type=["mp3", "wav"])
-
-   if uploaded_file:
-       file_size = len(uploaded_file.getvalue()) / (1024 * 1024)
-       st.write(f"File size: {file_size:.2f} MB")
-       
-       st.audio(uploaded_file, format='audio/wav')
-       
-       if st.button("Analyze Audio"):
-           if file_size > 200:
-               st.error("File size exceeds 200MB limit")
-           else:
-               results = process_audio(uploaded_file)
-               
-               if results:
-                   tab1, tab2, tab3 = st.tabs(["Speakers", "Transcription", "Summary"])
-                   
-                   with tab1:
-                       st.write("Speaker Timeline:")
-                       for turn, _, speaker in results["diarization"].itertracks(yield_label=True):
-                           col1, col2, col3 = st.columns([2,3,5])
-                           with col1:
-                               speaker_num = int(speaker.split('_')[1])
-                               colors = ['🔵', '🔴']
-                               st.write(f"{colors[speaker_num % 2]} {speaker}")
-                           with col2:
-                               st.write(f"{format_timestamp(turn.start)} → {format_timestamp(turn.end)}")
-                   
-                   with tab2:
-                       st.write("Transcription:")
-                       st.write(results["transcription"])
-                   
-                   with tab3:
-                       st.write("Summary:")
-                       st.write(results["summary"])
+    """Main application function."""
+    st.title("Multi-Speaker Audio Analyzer")
+    st.write("Upload an audio file (MP3/WAV) up to 5 minutes long for best performance")
+
+    uploaded_file = st.file_uploader("Choose a file", type=["mp3", "wav"])
+
+    if uploaded_file:
+        file_size = len(uploaded_file.getvalue()) / (1024 * 1024)
+        st.write(f"File size: {file_size:.2f} MB")
+        
+        st.audio(uploaded_file, format='audio/wav')
+        
+        if st.button("Analyze Audio"):
+            if file_size > 200:
+                st.error("File size exceeds 200MB limit")
+            else:
+                results = process_audio(uploaded_file)
+                
+                if results:
+                    tab1, tab2, tab3 = st.tabs(["Speakers", "Transcription", "Summary"])
+                    
+                    # Display speaker timeline
+                    with tab1:
+                        display_speaker_timeline(results)
+                    
+                    # Display transcription
+                    with tab2:
+                        display_transcription(results)
+                    
+                    # Display summary
+                    with tab3:
+                        display_summary(results)
+
+def display_speaker_timeline(results):
+    """Display speaker diarization results in a timeline format."""
+    st.write("Speaker Timeline:")
+    segments = TimeFormatter.format_speaker_segments(
+        results["diarization"], 
+        results["transcription"]
+    )
+    
+    if segments:
+        for segment in segments:
+            col1, col2, col3 = st.columns([2,3,5])
+            
+            with col1:
+                display_speaker_info(segment)
+            
+            with col2:
+                display_timestamp(segment)
+            
+            with col3:
+                display_text(segment)
+            
+            st.markdown("---")
+    else:
+        st.warning("No speaker segments detected")
+
+def display_speaker_info(segment):
+    """Display speaker information with color coding."""
+    speaker_num = int(segment['speaker'].split('_')[1])
+    colors = ['🔵', '🔴']
+    speaker_color = colors[speaker_num % len(colors)]
+    st.write(f"{speaker_color} {segment['speaker']}")
+
+def display_timestamp(segment):
+    """Display formatted timestamps."""
+    start_time = TimeFormatter.format_timestamp(segment['start'])
+    end_time = TimeFormatter.format_timestamp(segment['end'])
+    st.write(f"{start_time} → {end_time}")
+
+def display_text(segment):
+    """Display speaker's text."""
+    if segment['text']:
+        st.write(f"\"{segment['text']}\"")
+    else:
+        st.write("(no speech detected)")
 
 if __name__ == "__main__":
-   main()
+    main()