Update app.py

app.py

@@ -29,26 +29,27 @@ def load_model():
     return model, processor
 
 def chunk_audio(audio_file_path, chunk_length=30):
-    """Split audio into 30s mono chunks at 16kHz"""
+    """Split audio into chunks"""
     try:
+        # audio_file_path is already a string path when type="filepath"
+        audio = AudioSegment.from_file(audio_file_path)
+        
+        audios = []
         target_rate = 16000
         chunk_length_ms = chunk_length * 1000
         
-        # Load and resample
-        audio = AudioSegment.from_file(audio_file_path)
-        audio = audio.set_channels(1).set_frame_rate(target_rate)
-        
-        chunks = []
         for i in range(0, len(audio), chunk_length_ms):
             chunk = audio[i:i + chunk_length_ms]
+            chunk = chunk.set_channels(1).set_frame_rate(target_rate)
             
-            # Convert to NumPy array
-            samples = np.array(chunk.get_array_of_samples())
-            chunks.append(samples)
+            buf = io.BytesIO()
+            chunk.export(buf, format="wav")
+            buf.seek(0)
+            
+            rate, data = wavfile.read(buf)
+            audios.append(data)
         
-        print(f"Audio file split into {len(chunks)} chunks of ~{chunk_length}s")
-        return chunks
-    
+        return audios
     except Exception as e:
         raise gr.Error(f"Error processing audio file: {str(e)}. Make sure ffmpeg is installed.")
 
@@ -95,7 +96,6 @@ def search_audio(query, embeddings, audios, top_k=5):
     
     # Score against all embeddings
     scores = processor.score_multi_vector(query_embeddings, embeddings)
-    top_k = min(top_k, scores[0].shape[0])
     top_indices = scores[0].topk(top_k).indices.tolist()
     
     # Move model back to CPU
@@ -115,10 +115,10 @@ def audio_to_base64(data, rate=16000):
 def process_audio_rag(audio_file_path, query, chunk_length=30, use_openai=False, openai_key=None):
     """Main processing function"""
     if not audio_file_path:
-        return "Please upload an audio file", None
+        return "Please upload an audio file", None, None
     
     if not query:
-        return "Please enter a search query", None
+        return "Please enter a search query", None, None
     
     try:
         # Chunk audio
@@ -132,8 +132,7 @@ def process_audio_rag(audio_file_path, query, chunk_length=30, use_openai=False,
         
         # Prepare results
         result_text = f"Found {len(top_indices)} relevant audio chunks:\n"
-        result_text += f"Chunk indices: {top_indices}\n"
-        result_text += f"Total chunks in audio: {len(audios)}\n\n"
+        result_text += f"Chunk indices: {top_indices}\n\n"
         
         # Save first result as audio file
         first_chunk_path = "result_chunk.wav"
@@ -141,7 +140,6 @@ def process_audio_rag(audio_file_path, query, chunk_length=30, use_openai=False,
         
         # Optional: Use OpenAI for answer generation
         if use_openai and openai_key:
-            result_text += "Generating textual answer from retrieved audio chunks...\n\n"
             from openai import OpenAI
             client = OpenAI(api_key=openai_key)
             
@@ -164,14 +162,23 @@ def process_audio_rag(audio_file_path, query, chunk_length=30, use_openai=False,
                     model="gpt-4o-audio-preview",
                     messages=[{"role": "user", "content": content}]
                 )
-                result_text += f"OpenAI Answer: {completion.choices[0].message.content}"
+                result_text += f"\nOpenAI Answer: {completion.choices[0].message.content}"
             except Exception as e:
-                result_text += f"OpenAI Error: {str(e)}"
+                result_text += f"\nOpenAI Error: {str(e)}"
+        
+        # Create audio visualization
+        import matplotlib.pyplot as plt
+        fig, ax = plt.subplots(figsize=(10, 4))
+        ax.plot(audios[top_indices[0]])
+        ax.set_title(f"Waveform of top matching chunk (#{top_indices[0]})")
+        ax.set_xlabel("Samples")
+        ax.set_ylabel("Amplitude")
+        plt.tight_layout()
         
-        return result_text, first_chunk_path
+        return result_text, first_chunk_path, fig
         
     except Exception as e:
-        return f"Error: {str(e)}", None
+        return f"Error: {str(e)}", None, None
 
 # Create Gradio interface
 with gr.Blocks(title="AudioRAG Demo") as demo: