Update app.py

changed as per the trained model.

app.py

@@ -7,163 +7,126 @@ Date: January 2025
 """
 
 import streamlit as st
-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
+from pyannote.audio import Pipeline
+import whisper
+import tempfile
 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():
-    """
-    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
+   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
 
 def main():
-    """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)")
+   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"])
 
 if __name__ == "__main__":
-    main()
+   main()