Update app.py

app.py

@@ -18,9 +18,7 @@ try:
 
     # Try installing espeak with proper package manager commands
     try:
-        # Update package list first
         subprocess.run(['apt-get', 'update'], check=True)
-        # Try installing espeak first (more widely available)
         subprocess.run(['apt-get', 'install', '-y', 'espeak'], check=True)
     except subprocess.CalledProcessError:
         print("Warning: Could not install espeak. Attempting espeak-ng...")
@@ -33,7 +31,6 @@ except Exception as e:
     print(f"Warning: Initial setup error: {str(e)}")
     print("Continuing with limited functionality...")
 
-
 # --- Initialization (Do this ONCE) ---
 
 model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
@@ -66,7 +63,7 @@ try:
         from models import build_model  # type: ignore
         from kokoro import generate  # type: ignore
 
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'  # Correct device handling
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
         TTS_MODEL = build_model('Kokoro-82M/kokoro-v0_19.pth', device)
 
         # Load default voice
@@ -83,8 +80,6 @@ except Exception as e:
     print(f"Warning: Could not initialize Kokoro TTS: {str(e)}")
     TTS_ENABLED = False
 
-
-
 def get_web_results(query: str, max_results: int = 5) -> List[Dict[str, str]]:
     """Get web search results using DuckDuckGo"""
     try:
@@ -100,27 +95,19 @@ def get_web_results(query: str, max_results: int = 5) -> List[Dict[str, str]]:
         print(f"Error in web search: {e}")
         return []
 
-
 def format_prompt(query: str, context: List[Dict[str, str]]) -> str:
     """Format the prompt with web context"""
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     context_lines = '\n'.join([f'- [{res["title"]}]: {res["snippet"]}' for res in context])
     return f"""You are an intelligent search assistant. Answer the user's query using the provided web context.
-
 Current Time: {current_time}
-
 Important: For election-related queries, please distinguish clearly between different election years and types (presidential vs. non-presidential). Only use information from the provided web context.
-
 Query: {query}
-
 Web Context:
 {context_lines}
-
 Provide a detailed answer in markdown format. Include relevant information from sources and cite them using [1], [2], etc. If the query is about elections, clearly specify which year and type of election you're discussing.
-
 Answer:"""
 
-
 def format_sources(web_results: List[Dict[str, str]]) -> str:
     """Format sources with more details"""
     if not web_results:
@@ -143,7 +130,6 @@ def format_sources(web_results: List[Dict[str, str]]) -> str:
     sources_html += "</div>"
     return sources_html
 
-
 @spaces.GPU(duration=30)
 def generate_answer(prompt: str) -> str:
     """Generate answer using the DeepSeek model"""
@@ -168,47 +154,59 @@ def generate_answer(prompt: str) -> str:
     )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-
-
 @spaces.GPU(duration=30)
-def generate_speech_with_gpu(text: str, voice_name: str = 'af', tts_model = TTS_MODEL, voicepack = VOICEPACK) -> Tuple[int, np.ndarray] | None:
+def generate_speech_with_gpu(text: str, voice_name: str = 'af', tts_model=TTS_MODEL, voicepack=VOICEPACK) -> Tuple[int, np.ndarray] | None:
     """Generate speech from text using Kokoro TTS model."""
     if not TTS_ENABLED or tts_model is None:
         print("TTS is not enabled or model is not loaded.")
         return None
 
     try:
-        # Load voicepack if it hasn't been loaded or if a different voice is requested
-        if voice_name != 'af' or voicepack is None :
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
-            voicepack = torch.load(f'Kokoro-82M/voices/{voice_name}.pt', map_location=device, weights_only=True)
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        
+        # Handle voicepack loading
+        voice_file = f'Kokoro-82M/voices/{voice_name}.pt'
+        if voice_name == 'af' and voicepack is not None:
+            # Use the pre-loaded default voicepack
+            pass
+        elif os.path.exists(voice_file):
+            # Load the selected voicepack if it exists
+            voicepack = torch.load(voice_file, map_location=device, weights_only=True)
+        else:
+            # Fall back to default 'af' if selected voicepack is missing
+            print(f"Voicepack {voice_name}.pt not found. Falling back to default 'af'.")
+            voice_file = 'Kokoro-82M/voices/af.pt'
+            if os.path.exists(voice_file):
+                voicepack = torch.load(voice_file, map_location=device, weights_only=True)
+            else:
+                print("Default voicepack 'af.pt' not found. Cannot generate audio.")
+                return None
 
         # Clean the text
         clean_text = ' '.join([line for line in text.split('\n') if not line.startswith('#')])
         clean_text = clean_text.replace('[', '').replace(']', '').replace('*', '')
 
-        # Split long text into chunks (improved logic)
+        # Split long text into chunks
         max_chars = 1000
         chunks = []
         if len(clean_text) > max_chars:
             sentences = clean_text.split('.')
             current_chunk = ""
             for sentence in sentences:
-                if len(current_chunk) + len(sentence) + 1 < max_chars:  # +1 for the dot
+                if len(current_chunk) + len(sentence) + 1 < max_chars:
                     current_chunk += sentence + "."
                 else:
                     chunks.append(current_chunk.strip())
                     current_chunk = sentence + "."
-            if current_chunk:  # Add the last chunk
+            if current_chunk:
                 chunks.append(current_chunk.strip())
         else:
             chunks = [clean_text]
 
-
         # Generate audio for each chunk
         audio_chunks = []
         for chunk in chunks:
-            if chunk.strip():  # Only process non-empty chunks
+            if chunk.strip():
                 chunk_audio, _ = generate(tts_model, chunk, voicepack, lang='a')
                 if isinstance(chunk_audio, torch.Tensor):
                     chunk_audio = chunk_audio.cpu().numpy()
@@ -223,12 +221,8 @@ def generate_speech_with_gpu(text: str, voice_name: str = 'af', tts_model = TTS_
 
     except Exception as e:
         print(f"Error generating speech: {str(e)}")
-        import traceback
-        traceback.print_exc()
         return None
 
-
-
 def process_query(query: str, history: List[List[str]], selected_voice: str = 'af') -> Dict[str, Any]:
     """Process user query with streaming effect"""
     try:
@@ -242,11 +236,11 @@ def process_query(query: str, history: List[List[str]], selected_voice: str = 'a
         current_history = history + [[query, "*Searching...*"]]
 
         yield {
-            answer_output: gr.Markdown("*Searching & Thinking...*"),
-            sources_output: gr.HTML(sources_html),
-            search_btn: gr.Button("Searching...", interactive=False),
-            chat_history_display: current_history,
-            audio_output: None
+            'answer_output': gr.Markdown("*Searching & Thinking...*"),
+            'sources_output': gr.HTML(sources_html),
+            'search_btn': gr.Button("Searching...", interactive=False),
+            'chat_history_display': current_history,
+            'audio_output': None
         }
 
         # Generate answer
@@ -254,47 +248,47 @@ def process_query(query: str, history: List[List[str]], selected_voice: str = 'a
         answer = generate_answer(prompt)
         final_answer = answer.split("Answer:")[-1].strip()
 
-        # Update history *before* TTS (important for correct display)
+        # Update history before TTS
         updated_history = history + [[query, final_answer]]
 
-
         # Generate speech from the answer (only if enabled)
         if TTS_ENABLED:
-            yield {  # Intermediate update before TTS
-                answer_output: gr.Markdown(final_answer),
-                sources_output: gr.HTML(sources_html),
-                search_btn: gr.Button("Generating audio...", interactive=False),
-                chat_history_display: updated_history,
-                audio_output: None
+            yield {
+                'answer_output': gr.Markdown(final_answer),
+                'sources_output': gr.HTML(sources_html),
+                'search_btn': gr.Button("Generating audio...", interactive=False),
+                'chat_history_display': updated_history,
+                'audio_output': None
             }
             try:
                 audio = generate_speech_with_gpu(final_answer, selected_voice)
+                if audio is None:
+                    final_answer += "\n\n*Audio generation failed. The voicepack may be missing or incompatible.*"
             except Exception as e:
-                print(f"Error during TTS: {e}")
+                final_answer += f"\n\n*Error generating audio: {str(e)}*"
                 audio = None
         else:
+            final_answer += "\n\n*TTS is disabled. Audio not available.*"
             audio = None
 
-
-
         yield {
-            answer_output: gr.Markdown(final_answer),
-            sources_output: gr.HTML(sources_html),
-            search_btn: gr.Button("Search", interactive=True),
-            chat_history_display: updated_history,
-            audio_output: audio if audio is not None else gr.Audio(value=None)  # Ensure valid audio output
+            'answer_output': gr.Markdown(final_answer),
+            'sources_output': gr.HTML(sources_html),
+            'search_btn': gr.Button("Search", interactive=True),
+            'chat_history_display': updated_history,
+            'audio_output': audio if audio is not None else gr.Audio(value=None)
         }
 
     except Exception as e:
         error_message = str(e)
         if "GPU quota" in error_message:
-            error_message = "⚠️ GPU quota exceeded.  Please try again later when the daily quota resets."
+            error_message = "⚠️ GPU quota exceeded. Please try again later when the daily quota resets."
         yield {
-            answer_output: gr.Markdown(f"Error: {error_message}"),
-            sources_output: gr.HTML(sources_html), #Still show sources on error
-            search_btn: gr.Button("Search", interactive=True),
-            chat_history_display: history + [[query, f"*Error: {error_message}*"]],
-            audio_output: None
+            'answer_output': gr.Markdown(f"Error: {error_message}"),
+            'sources_output': gr.HTML(sources_html),
+            'search_btn': gr.Button("Search", interactive=True),
+            'chat_history_display': history + [[query, f"*Error: {error_message}*"]],
+            'audio_output': None
         }
 
 # Update the CSS for better contrast and readability
@@ -303,7 +297,6 @@ css = """
     max-width: 1200px !important;
     background-color: #f7f7f8 !important;
 }
-
 #header {
     text-align: center;
     margin-bottom: 2rem;
@@ -312,17 +305,14 @@ css = """
     border-radius: 12px;
     color: white;
 }
-
 #header h1 {
     color: white;
     font-size: 2.5rem;
     margin-bottom: 0.5rem;
 }
-
 #header h3 {
     color: #a8a9ab;
 }
-
 .search-container {
     background: #1a1b1e;
     border-radius: 12px;
@@ -330,40 +320,31 @@ css = """
     padding: 1rem;
     margin-bottom: 1rem;
 }
-
 .search-box {
     padding: 1rem;
     background: #2c2d30;
     border-radius: 8px;
     margin-bottom: 1rem;
 }
-
-/* Style the input textbox */
 .search-box input[type="text"] {
     background: #3a3b3e !important;
     border: 1px solid #4a4b4e !important;
     color: white !important;
     border-radius: 8px !important;
 }
-
 .search-box input[type="text"]::placeholder {
     color: #a8a9ab !important;
 }
-
-/* Style the search button */
 .search-box button {
     background: #2563eb !important;
     border: none !important;
 }
-
-/* Results area styling */
 .results-container {
     background: #2c2d30;
     border-radius: 8px;
     padding: 1rem;
     margin-top: 1rem;
 }
-
 .answer-box {
     background: #3a3b3e;
     border-radius: 8px;
@@ -371,19 +352,16 @@ css = """
     color: white;
     margin-bottom: 1rem;
 }
-
 .answer-box p {
     color: #e5e7eb;
     line-height: 1.6;
 }
-
 .sources-container {
     margin-top: 1rem;
     background: #2c2d30;
     border-radius: 8px;
     padding: 1rem;
 }
-
 .source-item {
     display: flex;
     padding: 12px;
@@ -392,21 +370,17 @@ css = """
     border-radius: 8px;
     transition: all 0.2s;
 }
-
 .source-item:hover {
     background: #4a4b4e;
 }
-
 .source-number {
     font-weight: bold;
     margin-right: 12px;
     color: #60a5fa;
 }
-
 .source-content {
     flex: 1;
 }
-
 .source-title {
     color: #60a5fa;
     font-weight: 500;
@@ -414,19 +388,16 @@ css = """
     display: block;
     margin-bottom: 4px;
 }
-
 .source-date {
     color: #a8a9ab;
     font-size: 0.9em;
     margin-left: 8px;
 }
-
 .source-snippet {
     color: #e5e7eb;
     font-size: 0.9em;
     line-height: 1.4;
 }
-
 .chat-history {
     max-height: 400px;
     overflow-y: auto;
@@ -435,47 +406,37 @@ css = """
     border-radius: 8px;
     margin-top: 1rem;
 }
-
 .examples-container {
     background: #2c2d30;
     border-radius: 8px;
     padding: 1rem;
     margin-top: 1rem;
 }
-
 .examples-container button {
     background: #3a3b3e !important;
     border: 1px solid #4a4b4e !important;
     color: #e5e7eb !important;
 }
-
-/* Markdown content styling */
 .markdown-content {
     color: #e5e7eb !important;
 }
-
 .markdown-content h1, .markdown-content h2, .markdown-content h3 {
     color: white !important;
 }
-
 .markdown-content a {
     color: #60a5fa !important;
 }
-
-/* Accordion styling */
 .accordion {
     background: #2c2d30 !important;
     border-radius: 8px !important;
     margin-top: 1rem !important;
 }
-
 .voice-selector {
     margin-top: 1rem;
     background: #2c2d30;
     border-radius: 8px;
     padding: 0.5rem;
 }
-
 .voice-selector select {
     background: #3a3b3e !important;
     color: white !important;
@@ -547,4 +508,4 @@ with gr.Blocks(title="AI Search Assistant", css=css, theme="dark") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True) 
+    demo.launch(share=True)