worable

app_qwen_tts.py

@@ -3,21 +3,19 @@ import torch
 import gradio as gr
 import numpy as np
 import requests
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from sentence_transformers import SentenceTransformer
 import base64
 import io
+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
 TOP_K = 3
-
-# FastAPI TTS endpoint
 TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
 
 # =========================================================
@@ -26,7 +24,7 @@ TTS_API_URL = "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
 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")
+    raise RuntimeError(f"{DOC_FILE} not found")
 
 # =========================================================
 # Load Qwen Model
@@ -41,20 +39,19 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 
 # =========================================================
-# Embedding Model
+# Embeddings
 # =========================================================
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
 # =========================================================
-# Document Chunking
+# 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]
-        chunks.append(" ".join(chunk))
+        chunks.append(" ".join(words[i:i+chunk_size]))
         i += chunk_size - overlap
     return chunks
 
@@ -65,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)
@@ -74,11 +71,11 @@ def retrieve_context(question, k=TOP_K):
     return "\n\n".join([DOC_CHUNKS[i] for i in top_ids])
 
 # =========================================================
-# Extract final answer
+# Extract answer
 # =========================================================
 def extract_final_answer(text: str) -> str:
     text = text.strip()
-    markers = ["assistant:", "assistant", "answer:", "final answer:"]
+    markers = ["assistant:", "answer:", "final answer:"]
     for m in markers:
         if m.lower() in text.lower():
             text = text.lower().split(m, 1)[-1].strip()
@@ -86,50 +83,36 @@ def extract_final_answer(text: str) -> str:
     return lines[-1] if lines else text
 
 # =========================================================
-# Qwen Inference
+# 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.\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.'"
+            "Answer ONLY the question in 1-2 sentences.\n"
+            "If not found, say '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)
 
 # =========================================================
-# TTS via API (returns NumPy audio)
+# TTS via API
 # =========================================================
 def tts_via_api(text: str):
     try:
-        payload = {"text": text}
-        resp = requests.post(TTS_API_URL, json=payload, timeout=60)
+        resp = requests.post(TTS_API_URL, json={"text": text}, timeout=60)
         resp.raise_for_status()
-        data = resp.json()
-        audio_b64 = data.get("audio", "")
+        audio_b64 = resp.json().get("audio", "")
         if not audio_b64:
             return None
-        # Decode base64 audio to bytes
         audio_bytes = base64.b64decode(audio_b64.split(",")[-1])
-        # Convert to np.float32
-        import soundfile as sf
         wav, sr = sf.read(io.BytesIO(audio_bytes), dtype='float32')
         return wav, sr
     except Exception as e:
@@ -137,26 +120,26 @@ def tts_via_api(text: str):
         return None
 
 # =========================================================
-# Gradio chat function
+# Chat function
 # =========================================================
 def chat(user_message, history):
     if not user_message.strip():
         return "", history
     try:
-        # Text answer
+        # 1️⃣ Text answer
         answer_text = answer_question(user_message)
 
-        # TTS
+        # 2️⃣ Audio
         tts_result = tts_via_api(answer_text)
         if tts_result is not None:
             wav, sr = tts_result
-            # Gradio can take NumPy array + sample rate directly
             audio_output = (sr, wav)
         else:
             audio_output = None
 
-        # Append tuple with text and playable audio
+        # 3️⃣ Append nicely formatted response
         history.append((user_message, [f"**Bot:** {answer_text}", audio_output]))
+
     except Exception as e:
         print(e)
         history.append((user_message, ["⚠️ Error generating answer or audio.", None]))
@@ -170,7 +153,8 @@ def reset_chat():
 # =========================================================
 def build_ui():
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("## 📄 Qwen Document Assistant + TTS")
+        gr.Markdown("# 📄 Qwen Document Assistant + TTS\nAsk a question and get a text + playable audio response.")
+
         chatbot = gr.Chatbot(height=450, type="tuples")
         msg = gr.Textbox(placeholder="Ask a question...", lines=2)
         send = gr.Button("Send")