Update app_qwen_tts.py

app_qwen_tts.py

@@ -1,12 +1,9 @@
 import os
-import io
-import base64
-import requests
 import torch
-import numpy as np
-import soundfile as sf
 import gradio as gr
-
+import numpy as np
+import base64
+import requests
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 
@@ -17,7 +14,7 @@ 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
+TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
 
 # =========================================================
 # Paths
@@ -31,9 +28,7 @@ if not os.path.exists(DOC_PATH):
 # =========================================================
 # Load Qwen Model
 # =========================================================
-print("🔄 Loading Qwen model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
@@ -41,7 +36,6 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 model.eval()
-print("✅ Qwen model loaded.")
 
 # =========================================================
 # Embedding Model
@@ -49,7 +43,7 @@ print("✅ Qwen model loaded.")
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
 # =========================================================
-# Document Chunking
+# Load Document
 # =========================================================
 def chunk_text(text, chunk_size=300, overlap=50):
     words = text.split()
@@ -68,7 +62,7 @@ DOC_CHUNKS = chunk_text(DOC_TEXT)
 DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, show_progress_bar=True)
 
 # =========================================================
-# Retrieval
+# Retrieve context
 # =========================================================
 def retrieve_context(question, k=TOP_K):
     q_emb = embedder.encode([question], normalize_embeddings=True)
@@ -89,20 +83,23 @@ def extract_final_answer(text: str) -> str:
     return lines[-1] if lines else text
 
 # =========================================================
-# Qwen Inference
+# Generate text 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"
-            "'I could not find this information in the document.'"
-        )},
+        {
+            "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}"}
     ]
 
@@ -110,74 +107,55 @@ def answer_question(question):
     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 via API (returns path to WAV for Gradio)
+# Call TTS API and get audio path
 # =========================================================
-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
-    }
+def tts_via_api(text: str):
     try:
+        payload = {"text": text}
         resp = requests.post(TTS_API_URL, json=payload, timeout=60)
         resp.raise_for_status()
         data = resp.json()
         audio_b64 = data.get("audio", "")
         if not audio_b64:
             return None
-
-        # Convert base64 to WAV for Gradio
+        audio_path = "/tmp/output.wav"
         audio_bytes = base64.b64decode(audio_b64)
-        audio_buffer = io.BytesIO(audio_bytes)
-        wav, sr = sf.read(audio_buffer, dtype="float32")
-        temp_path = "/tmp/response.wav"
-        sf.write(temp_path, wav, sr)
-        return temp_path
-
+        with open(audio_path, "wb") as f:
+            f.write(audio_bytes)
+        return audio_path
     except Exception as e:
         print(f"TTS API error: {e}")
         return None
 
 # =========================================================
-# Chat function for Gradio
+# Gradio Chat function
 # =========================================================
 def chat(user_message, history):
     if not user_message.strip():
         return "", history
 
     try:
-        # Generate text answer
+        # Generate text
         answer_text = answer_question(user_message)
 
         # Generate audio
         audio_path = tts_via_api(answer_text)
 
-        # Append user message
-        history.append({"role": "user", "content": user_message})
+        # Append tuple: (text, audio)
+        history.append((f"**Bot:** {answer_text}", audio_path))
 
-        # Append assistant message with text + audio
-        if audio_path:
-            history.append({"role": "assistant", "content": [f"**Bot:** {answer_text}", audio_path]})
-        else:
-            history.append({"role": "assistant", "content": f"**Bot:** {answer_text}"})
+        # Append user message
+        history.append((f"**You:** {user_message}", None))
 
     except Exception as e:
         print(e)
-        history.append({"role": "assistant", "content": "**⚠️ Error generating response.**"})
+        history.append(("⚠️ Error generating response", None))
 
     return "", history
 
@@ -185,14 +163,13 @@ def reset_chat():
     return []
 
 # =========================================================
-# Build Gradio UI
+# Build UI
 # =========================================================
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 📄 Qwen Document Assistant + TTS")
-        gr.Markdown("Ask questions and hear the answers as audio.")
+        gr.Markdown("# 📄 Qwen Document Assistant + TTS\nAsk questions and listen to answers.")
 
-        chatbot = gr.Chatbot(height=450, type="messages")
+        chatbot = gr.Chatbot(height=450, type="tuples")
         msg = gr.Textbox(placeholder="Ask a question...", lines=2)
         send = gr.Button("Send")
         clear = gr.Button("🧹 Clear")
@@ -201,7 +178,7 @@ def build_ui():
         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=True)
+        demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
 
 # =========================================================
 # Entrypoint