Update app_qwen_tts.py

app_qwen_tts.py

@@ -1,36 +1,37 @@
 import os
-import requests
 import torch
-import numpy as np
 import gradio as gr
+import numpy as np
+import soundfile as sf
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 
 # =========================================================
-# Config
+# 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"  # FastAPI TTS endpoint
 
 # =========================================================
-# Load document
+# 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")
 
-with open(DOC_PATH, "r", encoding="utf-8", errors="ignore") as f:
-    DOC_TEXT = f.read()
-
 # =========================================================
-# Qwen model
+# Load Qwen Model
 # =========================================================
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID, trust_remote_code=True
+)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
@@ -40,29 +41,41 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 
 # =========================================================
-# Embeddings for retrieval
+# Embedding Model
 # =========================================================
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
+# =========================================================
+# Document Chunking
+# =========================================================
 def chunk_text(text, chunk_size=300, overlap=50):
     words = text.split()
     chunks = []
     i = 0
     while i < len(words):
-        chunk = words[i:i+chunk_size]
+        chunk = words[i:i + chunk_size]
         chunks.append(" ".join(chunk))
         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)
 DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, show_progress_bar=True)
 
+# =========================================================
+# Retrieval
+# =========================================================
 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:"]
@@ -73,30 +86,46 @@ def extract_final_answer(text: str) -> str:
     return lines[-1] if lines else text
 
 # =========================================================
-# Qwen inference
+# Qwen Inference
 # =========================================================
-def answer_question(question: str) -> str:
+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"
-            "'I could not find this information in the document.'"
-        )},
-        {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}"}
+        {
+            "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"
+                "'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)
     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)
 
 # =========================================================
-# Call TTS API
+#  CPU-friendly TTS
+
 # =========================================================
 def tts_via_api(text: str, language_id="en", mode="Speak 🗣️", exaggeration=0.5, temperature=0.8, cfg_weight=0.5):
     payload = {
@@ -120,37 +149,33 @@ def tts_via_api(text: str, language_id="en", mode="Speak 🗣️", exaggeration=
         return None
 
 # =========================================================
-# Gradio Chat
+# Gradio Chatbot function
 # =========================================================
-def chat(user_message, history, language_id, mode, exaggeration, temperature, cfg_weight):
+def chat(user_message, history):
     if not user_message.strip():
         return "", history
+
     try:
-        # 1️⃣ Get Qwen answer
+        # 1️⃣ Generate text answer
         answer_text = answer_question(user_message)
 
-        # 2️⃣ Get TTS from API
-        audio_src = tts_via_api(answer_text, language_id, mode, exaggeration, temperature, cfg_weight)
+        # 2️⃣ Generate audio
+        sr, wav = tts_via_api(answer_text)
+
+        # Save temp wav for Gradio audio player
+        audio_path = "/tmp/output.wav"
+        import soundfile as sf
+        sf.write(audio_path, wav, sr)
 
-        # 3️⃣ Format bot message nicely with spacing
-        if audio_src:
-            # Use a small HTML wrapper for spacing
-            bot_content = [
-                f"<div style='margin-bottom:8px;'>{answer_text}</div>",  # text with margin
-                audio_src  # playable audio below
-            ]
-        else:
-            bot_content = [answer_text]
+        # 3️⃣ Append as tuple (text + audio)
+        history.append((user_message, [answer_text, audio_path]))
 
     except Exception as e:
         print(e)
-        bot_content = ["⚠️ Error generating answer or audio."]
+        history.append((user_message, ["⚠️ Error generating answer or audio."]))
 
-    # 4️⃣ Append as tuple: (user_message, bot_content)
-    history.append((user_message, bot_content))
     return "", history
 
-
 def reset_chat():
     return []
 
@@ -159,38 +184,22 @@ def reset_chat():
 # =========================================================
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("## 🧠 Qwen Document Chat + TTS (API-based)")
-
+        gr.Markdown("## 📄 Qwen Document Assistant + TTS")
         chatbot = gr.Chatbot(height=450, type="tuples")
-
-        with gr.Row():
-            msg = gr.Textbox(placeholder="Ask a question...", lines=2, scale=8)
-            send = gr.Button("🚀 Send", scale=2)
-
-        with gr.Row():
-            language_id = gr.Dropdown(["en","fr","hi","he"], value="en", label="TTS Language")
-            mode = gr.Radio(["Speak 🗣️", "Sing 🎵"], value="Speak 🗣️", label="TTS Mode")
-            exaggeration = gr.Slider(0.25, 2.0, step=0.05, value=0.5, label="Exaggeration")
-            temperature = gr.Slider(0.1, 2.0, step=0.05, value=0.8, label="Temperature")
-            cfg_weight = gr.Slider(0.2, 1.0, step=0.05, value=0.5, label="CFG / Pace")
-
+        msg = gr.Textbox(placeholder="Ask a question...", lines=2)
+        send = gr.Button("Send")
         clear = gr.Button("🧹 Clear")
 
-        send.click(
-            chat,
-            [msg, chatbot, language_id, mode, exaggeration, temperature, cfg_weight],
-            [msg, chatbot]
-        )
-        msg.submit(
-            chat,
-            [msg, chatbot, language_id, mode, exaggeration, temperature, cfg_weight],
-            [msg, chatbot]
-        )
+        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)
 
+# =========================================================
+# Entrypoint
+# =========================================================
 if __name__ == "__main__":
     print(f"✅ Loaded {len(DOC_CHUNKS)} chunks from {DOC_FILE}")
-    print(f"✅ Qwen Model: {MODEL_ID}")
+    print(f"✅ Model: {MODEL_ID}")
     build_ui()