Update app_qwen_tts.py

app_qwen_tts.py

@@ -1,31 +1,33 @@
 import os
+import requests
 import torch
 import gradio as gr
 import numpy as np
-import requests
-import base64
-import io
-import soundfile as sf
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 
 # =========================================================
 # Configuration
+# =========================================================
 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"
 
 # =========================================================
 # 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")
+    raise RuntimeError(f"❌ {DOC_FILE} not found next to app.py")
 
 # =========================================================
-# Load Qwen Model
+# Load Qwen model
+# =========================================================
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
@@ -36,17 +38,19 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 
 # =========================================================
-# Embeddings
+# Embedding Model for retrieval
+# =========================================================
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
 # =========================================================
-# Document chunking
+# Load document & chunk
+# =========================================================
 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]))
+        chunks.append(" ".join(words[i:i + chunk_size]))
         i += chunk_size - overlap
     return chunks
 
@@ -57,7 +61,8 @@ DOC_CHUNKS = chunk_text(DOC_TEXT)
 DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, show_progress_bar=True)
 
 # =========================================================
-# Retrieve context
+# Retrieval
+# =========================================================
 def retrieve_context(question, k=TOP_K):
     q_emb = embedder.encode([question], normalize_embeddings=True)
     scores = np.dot(DOC_EMBEDS, q_emb[0])
@@ -65,10 +70,11 @@ def retrieve_context(question, k=TOP_K):
     return "\n\n".join([DOC_CHUNKS[i] for i in top_ids])
 
 # =========================================================
-# Extract answer
+# Extract final answer
+# =========================================================
 def extract_final_answer(text: str) -> str:
     text = text.strip()
-    markers = ["assistant:", "answer:", "final answer:"]
+    markers = ["assistant:", "assistant", "answer:", "final answer:"]
     for m in markers:
         if m.lower() in text.lower():
             text = text.lower().split(m, 1)[-1].strip()
@@ -77,96 +83,110 @@ def extract_final_answer(text: str) -> str:
 
 # =========================================================
 # Qwen inference
+# =========================================================
 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 in 1-2 sentences.\n"
-            "If not found, say 'I could not find this information in the document.'"
+            "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)
+
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
     return extract_final_answer(decoded)
 
 # =========================================================
-# TTS via API
-def tts_via_api(text: str):
+# TTS via FastAPI
+# =========================================================
+def tts_via_api(text: str, language_id="en", mode="Speak 🗣️", exaggeration=0.5, temperature=0.8, cfg_weight=0.5):
+    payload = {
+        "text": text,
+        "language_id": language_id,
+        "mode": mode,
+        "exaggeration": exaggeration,
+        "temperature": temperature,
+        "cfg_weight": cfg_weight
+    }
     try:
-        resp = requests.post(TTS_API_URL, json={"text": text}, timeout=60)
+        resp = requests.post(TTS_API_URL, json=payload, timeout=60)
         resp.raise_for_status()
-        audio_b64 = resp.json().get("audio", "")
+        data = resp.json()
+        audio_b64 = data.get("audio", "")
         if not audio_b64:
             return None
-        audio_bytes = base64.b64decode(audio_b64.split(",")[-1])
-        wav, sr = sf.read(io.BytesIO(audio_bytes), dtype='float32')
-        return wav, sr
+        return f"data:audio/wav;base64,{audio_b64}"
     except Exception as e:
         print(f"TTS API error: {e}")
         return None
 
 # =========================================================
-# Chat function (text + audio separate boxes)
+# Chat function
+# =========================================================
 def chat(user_message, history):
     if not user_message.strip():
         return "", history
+
     try:
-        # 1️⃣ Text answer
+        # 1️⃣ Generate answer
         answer_text = answer_question(user_message)
 
-        # 2️⃣ Audio
-        tts_result = tts_via_api(answer_text)
-        if tts_result is not None:
-            wav, sr = tts_result
-            audio_output = (sr, wav)
-        else:
-            audio_output = None
-
-        # 3️⃣ Append as separate text + audio
-        history.append((user_message, answer_text, audio_output))
-
+        # 2️⃣ Generate audio
+        audio_data = tts_via_api(answer_text)
+
+        # 3️⃣ Append formatted message
+        history.append({
+            "role": "user",
+            "content": user_message
+        })
+        history.append({
+            "role": "assistant",
+            "content": [
+                gr.Markdown.update(value=f"**Bot:** {answer_text}"),
+                gr.Audio.update(value=audio_data, interactive=False) if audio_data else None
+            ]
+        })
     except Exception as e:
         print(e)
-        history.append((user_message, "⚠️ Error generating answer or audio.", None))
+        history.append({
+            "role": "assistant",
+            "content": "**⚠️ Error generating response.**"
+        })
+
     return "", history
 
 def reset_chat():
     return []
 
 # =========================================================
-# Gradio UI
+# Build UI
+# =========================================================
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 📄 Qwen Document Assistant + TTS\nAsk a question and get a text + playable audio response.")
-
-        chatbot = gr.Chatbot(height=500, type="messages")  # 'messages' so we can use custom formatting
-
+        gr.Markdown("## 📄 Qwen Document Assistant + TTS\nAsk questions and listen to answers!")
+        chatbot = gr.Chatbot(height=500, type="messages")
         msg = gr.Textbox(placeholder="Ask a question...", lines=2)
         send = gr.Button("Send")
         clear = gr.Button("🧹 Clear")
 
-        def format_history(history):
-            formatted = []
-            for user_msg, bot_text, bot_audio in history:
-                formatted.append([f"**You:** {user_msg}", None])
-                formatted.append([f"**Bot:** {bot_text}", bot_audio])
-            return formatted
-
-        def chat_with_format(msg_input, history):
-            _, history = chat(msg_input, history)
-            return "", format_history(history)
-
-        send.click(chat_with_format, [msg, chatbot], [msg, chatbot])
-        msg.submit(chat_with_format, [msg, chatbot], [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)
+        demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
 
+# =========================================================
+# Entrypoint
 # =========================================================
 if __name__ == "__main__":
     print(f"✅ Loaded {len(DOC_CHUNKS)} chunks from {DOC_FILE}")