Update app.py

app.py

@@ -11,18 +11,12 @@ import io
 @st.cache_resource
 def load_models():
     try:
-        # Updated to 3.1 with parameters
         diarization = Pipeline.from_pretrained(
             "pyannote/speaker-diarization-3.1",
             use_auth_token=st.secrets["hf_token"]
-        ).instantiate({
-            "onset": 0.3,
-            "offset": 0.3,
-            "min_duration_on": 0.1,
-            "min_duration_off": 0.1
-        })
+        )
         
-        transcriber = whisper.load_model("base")
+        transcriber = whisper.load_model("small")
         
         summarizer = tf_pipeline(
             "summarization",
@@ -78,7 +72,7 @@ def process_audio(audio_file, max_duration=600):
             
             return {
                 "diarization": diarization_result,
-                "transcription": transcription,  # Return full transcription object
+                "transcription": transcription,
                 "summary": summary[0]["summary_text"]
             }
             
@@ -86,26 +80,24 @@ def process_audio(audio_file, max_duration=600):
         st.error(f"Error processing audio: {str(e)}")
         return None
 
-def format_speaker_segments(diarization_result, transcription):
+def format_speaker_segments(diarization_result):
     formatted_segments = []
-    audio_duration = transcription.get('duration', 0)
     
     for turn, _, speaker in diarization_result.itertracks(yield_label=True):
-        # Skip invalid timestamps
-        if turn.start > audio_duration or turn.end > audio_duration:
-            continue
-            
-        # Only add segments with meaningful duration
-        if (turn.end - turn.start) >= 0.1:  # 100ms minimum
+        if turn.start is not None and turn.end is not None:
             formatted_segments.append({
                 'speaker': speaker,
-                'start': turn.start,
-                'end': turn.end,
-                'duration': turn.end - turn.start
+                'start': float(turn.start),
+                'end': float(turn.end)
             })
     
     return formatted_segments
 
+def format_timestamp(seconds):
+    minutes = int(seconds // 60)
+    seconds = seconds % 60
+    return f"{minutes:02d}:{seconds:05.2f}"
+
 def main():
     st.title("Multi-Speaker Audio Analyzer")
     st.write("Upload an audio file (MP3/WAV) up to 5 minutes long for best performance")
@@ -129,30 +121,21 @@ def main():
                     
                     with tab1:
                         st.write("Speaker Timeline:")
+                        segments = format_speaker_segments(results["diarization"])
                         
-                        segments = format_speaker_segments(
-                            results["diarization"], 
-                            results["transcription"]
-                        )
-                        
-                        # Display segments with proper time formatting
                         for segment in segments:
                             col1, col2 = st.columns([2,8])
                             
                             with col1:
                                 speaker_num = int(segment['speaker'].split('_')[1])
-                                colors = ['🔵', '🔴']  # Simplified to two colors
+                                colors = ['🔵', '🔴']  # Two colors for alternating speakers
                                 speaker_color = colors[speaker_num % len(colors)]
                                 st.write(f"{speaker_color} {segment['speaker']}")
                             
                             with col2:
-                                mm_start = int(segment['start'] // 60)
-                                ss_start = segment['start'] % 60
-                                mm_end = int(segment['end'] // 60)
-                                ss_end = segment['end'] % 60
-                                
-                                time_str = f"{mm_start:02d}:{ss_start:05.2f} → {mm_end:02d}:{ss_end:05.2f}"
-                                st.write(time_str)
+                                start_time = format_timestamp(segment['start'])
+                                end_time = format_timestamp(segment['end'])
+                                st.write(f"{start_time} → {end_time}")
                             
                             st.markdown("---")