Update app_qwen_tts.py

app_qwen_tts.py

@@ -1,56 +1,51 @@
 import os
 import io
 import base64
-import time
+import requests
 import torch
 import gradio as gr
 import numpy as np
-import soundfile as sf
-import requests
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
+from scipy.io.wavfile import write as write_wav
 
-# =======================
-# Configuration
-# =======================
+# =====================================================
+# CONFIG
+# =====================================================
 MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 DOC_FILE = "general.md"
+TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
 MAX_NEW_TOKENS = 200
 TOP_K = 3
-TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"  # your FastAPI TTS endpoint
 
-# =======================
-# Load document
-# =======================
+# =====================================================
+# 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")
-
 with open(DOC_PATH, "r", encoding="utf-8", errors="ignore") as f:
     DOC_TEXT = f.read()
 
-# =======================
-# Chunk document
-# =======================
+# =====================================================
+# CHUNK + EMBED
+# =====================================================
 def chunk_text(text, chunk_size=300, overlap=50):
     words = text.split()
-    chunks = []
-    i = 0
+    chunks, i = [], 0
     while i < len(words):
-        chunk = words[i:i + chunk_size]
-        chunks.append(" ".join(chunk))
+        chunks.append(" ".join(words[i:i + chunk_size]))
         i += chunk_size - overlap
     return chunks
 
 DOC_CHUNKS = chunk_text(DOC_TEXT)
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True)
 
-# =======================
-# Load models
-# =======================
+# =====================================================
+# LOAD QWEN
+# =====================================================
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
@@ -59,108 +54,107 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
-DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, show_progress_bar=True)
-
-# =======================
-# Utilities
-# =======================
-def retrieve_context(question, k=TOP_K):
+# =====================================================
+# RETRIEVAL
+# =====================================================
+def retrieve_context(question):
     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])
-
-def extract_final_answer(text: str) -> str:
-    text = text.strip()
-    markers = ["assistant:", "assistant", "answer:", "final answer:"]
-    for m in markers:
-        if m.lower() in text.lower():
-            text = text.lower().split(m, 1)[-1].strip()
-    lines = [l.strip() for l in text.split("\n") if l.strip()]
-    return lines[-1] if lines else text
-
-# =======================
-# Qwen inference
-# =======================
-def answer_question(question: str) -> str:
+    top_ids = scores.argsort()[-TOP_K:][::-1]
+    return "\n\n".join(DOC_CHUNKS[i] for i in top_ids)
+
+# =====================================================
+# QWEN ANSWER
+# =====================================================
+def answer_question(question):
     context = retrieve_context(question)
+
     messages = [
         {
             "role": "system",
             "content": (
                 "You are a strict document-based Q&A assistant.\n"
                 "Answer ONLY the question.\n"
-                "Do NOT repeat the context or the question.\n"
-                "Respond in 1–2 sentences.\n"
-                "If the answer is not present, say:\n"
+                "Do NOT repeat context.\n"
+                "Respond in 1 sentence.\n"
+                "If not found, say:\n"
                 "'I could not find this information in the document.'"
             )
         },
         {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}"}
     ]
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    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)
 
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    return extract_final_answer(decoded)
-
-# =======================
-# TTS via FastAPI
-# =======================
-def generate_tts_base64(text: str, language_id="en") -> str:
-    try:
-        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()
-        audio_b64 = resp.json().get("audio", "")
-        return audio_b64
-    except Exception as e:
-        print(f"TTS error: {e}")
-        return None
-
-# =======================
-# Chat function for Gradio
-# =======================
-def chat(user_message, history):
-    if not user_message.strip():
-        return "", history
-
-    # 1️⃣ Text answer immediately
-    answer_text = answer_question(user_message)
-    history.append((user_message, [answer_text, None]))  # audio placeholder
-
-    # 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
-# =======================
+    return decoded.split("\n")[-1].strip()
+
+# =====================================================
+# TTS (BASE64 → WAV)
+# =====================================================
+def generate_audio(text):
+    payload = {"text": text, "language_id": "en", "mode": "Speak 🗣️"}
+    r = requests.post(TTS_API_URL, json=payload, timeout=None)
+    audio_b64 = r.json()["audio"]
+
+    audio_bytes = base64.b64decode(audio_b64)
+    wav_path = "/tmp/output.wav"
+
+    with open(wav_path, "wb") as f:
+        f.write(audio_bytes)
+
+    return wav_path
+
+# =====================================================
+# MAIN HANDLER
+# =====================================================
+def run_pipeline(question):
+    if not question.strip():
+        return "", None
+
+    # 1️⃣ TEXT FIRST
+    answer = answer_question(question)
+
+    # 2️⃣ AUDIO (SLOW, NO TIMEOUT)
+    audio_path = generate_audio(answer)
+
+    return answer, audio_path
+
+# =====================================================
+# 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)
+    gr.Markdown("## 🤖 Document Q&A with Voice")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            user_input = gr.Textbox(
+                label="Your Question",
+                placeholder="Who is CEO of OhamLab?",
+                lines=4
+            )
+            ask_btn = gr.Button("Ask")
+
+        with gr.Column(scale=1):
+            answer_text = gr.Markdown(
+                label="Assistant Answer",
+                value="**Bot:** _Waiting for question..._"
+            )
+            answer_audio = gr.Audio(
+                label="Assistant Voice",
+                type="filepath"
+            )
+
+    ask_btn.click(
+        fn=run_pipeline,
+        inputs=user_input,
+        outputs=[answer_text, answer_audio]
+    )
 
 demo.launch(server_name="0.0.0.0", server_port=7860, share=False)