Update app_qwen_tts.py

app_qwen_tts.py

@@ -1,8 +1,12 @@
 import os
+import io
+import base64
 import requests
 import torch
-import gradio as gr
 import numpy as np
+import soundfile as sf
+import gradio as gr
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 
@@ -13,22 +17,23 @@ MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 DOC_FILE = "general.md"
 MAX_NEW_TOKENS = 200
 TOP_K = 3
-
-# Your TTS FastAPI endpoint
-TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
+TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"  # FastAPI TTS endpoint
 
 # =========================================================
 # Paths
 # =========================================================
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 DOC_PATH = os.path.join(BASE_DIR, DOC_FILE)
+
 if not os.path.exists(DOC_PATH):
     raise RuntimeError(f"❌ {DOC_FILE} not found next to app.py")
 
 # =========================================================
-# Load Qwen model
+# Load Qwen Model
 # =========================================================
+print("🔄 Loading Qwen model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
@@ -36,21 +41,23 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 model.eval()
+print("✅ Qwen model loaded.")
 
 # =========================================================
-# Embedding Model for retrieval
+# Embedding Model
 # =========================================================
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
 # =========================================================
-# Load document & chunk
+# Document Chunking
 # =========================================================
 def chunk_text(text, chunk_size=300, overlap=50):
     words = text.split()
     chunks = []
     i = 0
     while i < len(words):
-        chunks.append(" ".join(words[i:i + chunk_size]))
+        chunk = words[i:i + chunk_size]
+        chunks.append(" ".join(chunk))
         i += chunk_size - overlap
     return chunks
 
@@ -82,10 +89,11 @@ def extract_final_answer(text: str) -> str:
     return lines[-1] if lines else text
 
 # =========================================================
-# Qwen inference
+# Qwen Inference
 # =========================================================
 def answer_question(question):
     context = retrieve_context(question)
+
     messages = [
         {"role": "system", "content": (
             "You are a strict document-based Q&A assistant.\n"
@@ -97,17 +105,23 @@ def answer_question(question):
         )},
         {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}"}
     ]
+
     prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
-        output = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, temperature=0.3, do_sample=True)
+        output = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            temperature=0.3,
+            do_sample=True
+        )
 
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
     return extract_final_answer(decoded)
 
 # =========================================================
-# TTS via FastAPI
+# TTS via API (returns path to WAV for Gradio)
 # =========================================================
 def tts_via_api(text: str, language_id="en", mode="Speak 🗣️", exaggeration=0.5, temperature=0.8, cfg_weight=0.5):
     payload = {
@@ -125,61 +139,60 @@ def tts_via_api(text: str, language_id="en", mode="Speak 🗣️", exaggeration=
         audio_b64 = data.get("audio", "")
         if not audio_b64:
             return None
-        return f"data:audio/wav;base64,{audio_b64}"
+
+        # Convert base64 to WAV for Gradio
+        audio_bytes = base64.b64decode(audio_b64)
+        audio_buffer = io.BytesIO(audio_bytes)
+        wav, sr = sf.read(audio_buffer, dtype="float32")
+        temp_path = "/tmp/response.wav"
+        sf.write(temp_path, wav, sr)
+        return temp_path
+
     except Exception as e:
         print(f"TTS API error: {e}")
         return None
 
 # =========================================================
-# Chat function
-# =========================================================
-# =========================================================
-# Chat function
+# Chat function for Gradio
 # =========================================================
 def chat(user_message, history):
     if not user_message.strip():
         return "", history
 
     try:
-        # 1️⃣ Generate text answer
+        # Generate text answer
         answer_text = answer_question(user_message)
 
-        # 2️⃣ Generate audio
-        audio_data = tts_via_api(answer_text)
+        # Generate audio
+        audio_path = tts_via_api(answer_text)
 
-        # 3️⃣ Append messages in 'messages' format
+        # Append user message
         history.append({"role": "user", "content": user_message})
-        if audio_data:
-            history.append({
-                "role": "assistant",
-                "content": [f"**Bot:** {answer_text}", audio_data]
-            })
+
+        # Append assistant message with text + audio
+        if audio_path:
+            history.append({"role": "assistant", "content": [f"**Bot:** {answer_text}", audio_path]})
         else:
-            history.append({
-                "role": "assistant",
-                "content": f"**Bot:** {answer_text}"
-            })
+            history.append({"role": "assistant", "content": f"**Bot:** {answer_text}"})
 
     except Exception as e:
         print(e)
-        history.append({
-            "role": "assistant",
-            "content": "**⚠️ Error generating response.**"
-        })
+        history.append({"role": "assistant", "content": "**⚠️ Error generating response.**"})
 
     return "", history
 
-
 def reset_chat():
     return []
 
 # =========================================================
-# Build UI
+# Build Gradio UI
 # =========================================================
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("## 📄 Qwen Document Assistant + TTS\nAsk questions and listen to answers!")
-        chatbot = gr.Chatbot(height=500, type="messages")
+        gr.Markdown("# 📄 Qwen Document Assistant + TTS")
+        gr.Markdown("Ask questions and hear the answers as audio.")
+
+        chatbot = gr.Chatbot(height=450, type="messages")
         msg = gr.Textbox(placeholder="Ask a question...", lines=2)
         send = gr.Button("Send")
         clear = gr.Button("🧹 Clear")