Update app_qwen_tts.py

app_qwen_tts.py

@@ -5,13 +5,12 @@ import numpy as np
 import requests
 import base64
 import io
+import soundfile as sf
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
-import soundfile as sf
 
 # =========================================================
 # Configuration
-# =========================================================
 MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 DOC_FILE = "general.md"
 MAX_NEW_TOKENS = 200
@@ -20,7 +19,6 @@ 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):
@@ -28,7 +26,6 @@ if not os.path.exists(DOC_PATH):
 
 # =========================================================
 # Load Qwen Model
-# =========================================================
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
@@ -40,12 +37,10 @@ model.eval()
 
 # =========================================================
 # Embeddings
-# =========================================================
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
 # =========================================================
 # Document chunking
-# =========================================================
 def chunk_text(text, chunk_size=300, overlap=50):
     words = text.split()
     chunks = []
@@ -63,7 +58,6 @@ DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, show_progres
 
 # =========================================================
 # Retrieve context
-# =========================================================
 def retrieve_context(question, k=TOP_K):
     q_emb = embedder.encode([question], normalize_embeddings=True)
     scores = np.dot(DOC_EMBEDS, q_emb[0])
@@ -72,7 +66,6 @@ def retrieve_context(question, k=TOP_K):
 
 # =========================================================
 # Extract answer
-# =========================================================
 def extract_final_answer(text: str) -> str:
     text = text.strip()
     markers = ["assistant:", "answer:", "final answer:"]
@@ -84,7 +77,6 @@ def extract_final_answer(text: str) -> str:
 
 # =========================================================
 # Qwen inference
-# =========================================================
 def answer_question(question):
     context = retrieve_context(question)
     messages = [
@@ -104,7 +96,6 @@ def answer_question(question):
 
 # =========================================================
 # TTS via API
-# =========================================================
 def tts_via_api(text: str):
     try:
         resp = requests.post(TTS_API_URL, json={"text": text}, timeout=60)
@@ -120,8 +111,7 @@ def tts_via_api(text: str):
         return None
 
 # =========================================================
-# Chat function
-# =========================================================
+# Chat function (text + audio separate boxes)
 def chat(user_message, history):
     if not user_message.strip():
         return "", history
@@ -137,12 +127,12 @@ def chat(user_message, history):
         else:
             audio_output = None
 
-        # 3️⃣ Append nicely formatted response
-        history.append((user_message, [f"**Bot:** {answer_text}", audio_output]))
+        # 3️⃣ Append as separate text + audio
+        history.append((user_message, answer_text, audio_output))
 
     except Exception as e:
         print(e)
-        history.append((user_message, ["⚠️ Error generating answer or audio.", None]))
+        history.append((user_message, "⚠️ Error generating answer or audio.", None))
     return "", history
 
 def reset_chat():
@@ -150,24 +140,33 @@ def reset_chat():
 
 # =========================================================
 # Gradio 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=450, type="tuples")
+        chatbot = gr.Chatbot(height=500, type="messages")  # 'messages' so we can use custom formatting
+
         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])
+        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])
         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}")