Update app_qwen_tts.py

app_qwen_tts.py

@@ -1,50 +1,39 @@
 import os
+import io
+import base64
+import time
 import torch
 import gradio as gr
 import numpy as np
-import base64
+import soundfile as sf
 import requests
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
-import asyncio
 
-# =========================================================
+# =======================
 # Configuration
-# =========================================================
+# =======================
 MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 DOC_FILE = "general.md"
 MAX_NEW_TOKENS = 200
 TOP_K = 3
-TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
+TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"  # your FastAPI TTS endpoint
 
-# =========================================================
-# Paths
-# =========================================================
+# =======================
+# Load document
+# =======================
 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
-# =========================================================
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    trust_remote_code=True
-)
-model.eval()
+if not os.path.exists(DOC_PATH):
+    raise RuntimeError(f"{DOC_FILE} not found next to app.py")
 
-# =========================================================
-# Embedding Model
-# =========================================================
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
+with open(DOC_PATH, "r", encoding="utf-8", errors="ignore") as f:
+    DOC_TEXT = f.read()
 
-# =========================================================
-# Load Document
-# =========================================================
+# =======================
+# Chunk document
+# =======================
 def chunk_text(text, chunk_size=300, overlap=50):
     words = text.split()
     chunks = []
@@ -55,24 +44,33 @@ def chunk_text(text, chunk_size=300, overlap=50):
         i += chunk_size - overlap
     return chunks
 
-with open(DOC_PATH, "r", encoding="utf-8", errors="ignore") as f:
-    DOC_TEXT = f.read()
-
 DOC_CHUNKS = chunk_text(DOC_TEXT)
+
+# =======================
+# Load models
+# =======================
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    trust_remote_code=True
+)
+model.eval()
+
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
 DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, show_progress_bar=True)
 
-# =========================================================
-# Retrieval
-# =========================================================
+# =======================
+# Utilities
+# =======================
 def retrieve_context(question, k=TOP_K):
     q_emb = embedder.encode([question], normalize_embeddings=True)
     scores = np.dot(DOC_EMBEDS, q_emb[0])
     top_ids = scores.argsort()[-k:][::-1]
     return "\n\n".join([DOC_CHUNKS[i] for i in top_ids])
 
-# =========================================================
-# Extract final answer
-# =========================================================
 def extract_final_answer(text: str) -> str:
     text = text.strip()
     markers = ["assistant:", "assistant", "answer:", "final answer:"]
@@ -82,12 +80,11 @@ def extract_final_answer(text: str) -> str:
     lines = [l.strip() for l in text.split("\n") if l.strip()]
     return lines[-1] if lines else text
 
-# =========================================================
-# Generate text
-# =========================================================
-def answer_question(question):
+# =======================
+# Qwen inference
+# =======================
+def answer_question(question: str) -> str:
     context = retrieve_context(question)
-
     messages = [
         {
             "role": "system",
@@ -102,93 +99,68 @@ 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 API call
-# =========================================================
-def tts_via_api(text: str):
+# =======================
+# TTS via FastAPI
+# =======================
+def generate_tts_base64(text: str, language_id="en") -> str:
     try:
-        payload = {"text": text}
-        resp = requests.post(TTS_API_URL, json=payload)
+        payload = {"text": text, "language_id": language_id, "mode": "Speak 🗣️"}
+        resp = requests.post(TTS_API_URL, json=payload, timeout=None)  # no timeout
         resp.raise_for_status()
-        data = resp.json()
-        audio_b64 = data.get("audio", "")
-        if not audio_b64:
-            return None
-        audio_path = "/tmp/output.wav"
-        audio_bytes = base64.b64decode(audio_b64)
-        with open(audio_path, "wb") as f:
-            f.write(audio_bytes)
-        return audio_path
+        audio_b64 = resp.json().get("audio", "")
+        return audio_b64
     except Exception as e:
-        print(f"TTS API error: {e}")
+        print(f"TTS error: {e}")
         return None
 
-# =========================================================
-# Async chat
-# =========================================================
-async def chat_async(user_message, history):
+# =======================
+# Chat function for Gradio
+# =======================
+def chat(user_message, history):
     if not user_message.strip():
         return "", history
 
-    # Show user message immediately
-    history.append((user_message, None))
+    # 1️⃣ Text answer immediately
+    answer_text = answer_question(user_message)
+    history.append((user_message, [answer_text, None]))  # audio placeholder
 
-    try:
-        # 1️⃣ Generate text answer immediately
-        answer_text = answer_question(user_message)
-
-        # Append answer text only (audio pending)
-        history.append((None, (answer_text, None)))
-
-        # 2️⃣ Generate audio in background
-        audio_path = await asyncio.to_thread(tts_via_api, answer_text)
-
-        # Update the last bot message with audio
-        history[-1] = (None, (answer_text, audio_path))
-    except Exception as e:
-        print(e)
-        history.append((None, ("⚠️ Error generating response", None)))
+    # 2️⃣ Generate audio asynchronously
+    audio_b64 = generate_tts_base64(answer_text)
+    if audio_b64:
+        history[-1][1][1] = f"data:audio/wav;base64,{audio_b64}"
 
     return "", history
 
 def reset_chat():
     return []
 
-# =========================================================
+# =======================
 # Build UI
-# =========================================================
-def build_ui():
-    with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 📄 Qwen Document Assistant + TTS (Async Audio)")
-        gr.Markdown("Ask a question and hear the answer. Text appears immediately, audio may take a few minutes.")
-
-        # Chat history: left=User, right=Bot
-        chatbot = gr.Chatbot(height=450, type="tuples")
-
-        msg = gr.Textbox(placeholder="Ask a question...", lines=2)
-        send = gr.Button("Send")
-        clear = gr.Button("🧹 Clear")
-
-        send.click(chat_async, [msg, chatbot], [msg, chatbot])
-        msg.submit(chat_async, [msg, chatbot], [msg, chatbot])
-        clear.click(reset_chat, outputs=chatbot)
-
-        demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
-
-# =========================================================
-# Entrypoint
-# =========================================================
-if __name__ == "__main__":
-    print(f"✅ Loaded {len(DOC_CHUNKS)} chunks from {DOC_FILE}")
-    print(f"✅ Model: {MODEL_ID}")
-    build_ui()
+# =======================
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("## 📄 Qwen Document Assistant + TTS\nText appears instantly; audio plays once ready.")
+    
+    chatbot = gr.Chatbot(height=450, type="tuples")
+    msg = gr.Textbox(placeholder="Ask a question...", lines=2)
+    send = gr.Button("Send")
+    clear = gr.Button("Clear")
+
+    send.click(chat, [msg, chatbot], [msg, chatbot])
+    msg.submit(chat, [msg, chatbot], [msg, chatbot])
+    clear.click(reset_chat, outputs=chatbot)
+
+demo.launch(server_name="0.0.0.0", server_port=7860, share=False)