multi-lang support test

app_qwen_tts_fast.py

@@ -13,27 +13,32 @@ from sentence_transformers import SentenceTransformer
 # CONFIG
 # =====================================================
 MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
-DOC_FILE = "general.md"
+DOC_FILE_EN = "general.md"
+DOC_FILE_HI = "general-hi.md"
 TTS_API_URL = os.getenv(
     "TTS_API_URL",
     "https://rahul7star-Chatterbox-Multilingual-TTS-API.hf.space/tts"
 )
 MAX_NEW_TOKENS = 128
 TOP_K = 3
-
 SESSION = requests.Session()
 
 # =====================================================
-# LOAD DOCUMENT
+# LOAD DOCUMENTS
 # =====================================================
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-DOC_PATH = os.path.join(BASE_DIR, DOC_FILE)
+DOC_PATH_EN = os.path.join(BASE_DIR, DOC_FILE_EN)
+DOC_PATH_HI = os.path.join(BASE_DIR, DOC_FILE_HI)
+
+for path, name in [(DOC_PATH_EN, DOC_FILE_EN), (DOC_PATH_HI, DOC_FILE_HI)]:
+    if not os.path.exists(path):
+        raise RuntimeError(f"{name} not found")
 
-if not os.path.exists(DOC_PATH):
-    raise RuntimeError(f"{DOC_FILE} not found")
+with open(DOC_PATH_EN, "r", encoding="utf-8", errors="ignore") as f:
+    DOC_TEXT_EN = f.read()
 
-with open(DOC_PATH, "r", encoding="utf-8", errors="ignore") as f:
-    DOC_TEXT = f.read()
+with open(DOC_PATH_HI, "r", encoding="utf-8", errors="ignore") as f:
+    DOC_TEXT_HI = f.read()
 
 # =====================================================
 # CHUNK + EMBED
@@ -46,10 +51,12 @@ def chunk_text(text, chunk_size=300, overlap=50):
         i += chunk_size - overlap
     return chunks
 
-DOC_CHUNKS = chunk_text(DOC_TEXT)
+DOC_CHUNKS_EN = chunk_text(DOC_TEXT_EN)
+DOC_CHUNKS_HI = chunk_text(DOC_TEXT_HI)
 
 embedder = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
-DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, batch_size=32)
+DOC_EMBEDS_EN = embedder.encode(DOC_CHUNKS_EN, normalize_embeddings=True, batch_size=32)
+DOC_EMBEDS_HI = embedder.encode(DOC_CHUNKS_HI, normalize_embeddings=True, batch_size=32)
 
 # =====================================================
 # LOAD QWEN MODEL (CPU only)
@@ -57,7 +64,7 @@ DOC_EMBEDS = embedder.encode(DOC_CHUNKS, normalize_embeddings=True, batch_size=3
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    device_map="cpu",  # strictly CPU
+    device_map="cpu",
     torch_dtype=torch.float32,
     trust_remote_code=True
 )
@@ -67,17 +74,22 @@ model.eval()
 # RETRIEVAL WITH CACHE
 # =====================================================
 @lru_cache(maxsize=256)
-def retrieve_context(question: str):
+def retrieve_context(question: str, lang: str):
     q_emb = embedder.encode([question], normalize_embeddings=True)
-    scores = np.dot(DOC_EMBEDS, q_emb[0])
-    top_ids = scores.argsort()[-TOP_K:][::-1]
-    return "\n\n".join(DOC_CHUNKS[i] for i in top_ids)
+    if lang == "hi":
+        scores = np.dot(DOC_EMBEDS_HI, q_emb[0])
+        top_ids = scores.argsort()[-TOP_K:][::-1]
+        return "\n\n".join(DOC_CHUNKS_HI[i] for i in top_ids)
+    else:
+        scores = np.dot(DOC_EMBEDS_EN, q_emb[0])
+        top_ids = scores.argsort()[-TOP_K:][::-1]
+        return "\n\n".join(DOC_CHUNKS_EN[i] for i in top_ids)
 
 # =====================================================
 # QWEN ANSWER (CPU optimized)
 # =====================================================
-def answer_question(question: str) -> str:
-    context = retrieve_context(question)
+def answer_question(question: str, lang: str = "en") -> str:
+    context = retrieve_context(question, lang)
 
     messages = [
         {
@@ -90,16 +102,10 @@ def answer_question(question: str) -> str:
                 "'I could not find this information in the document.'"
             )
         },
-        {
-            "role": "user",
-            "content": f"Context:\n{context}\n\nQuestion:\n{question}"
-        }
+        {"role": "user", "content": f"Context:\n{context}\n\nQuestion:\n{question}"}
     ]
 
-    prompt = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
 
     with torch.no_grad():
@@ -118,24 +124,19 @@ def answer_question(question: str) -> str:
 # =====================================================
 @lru_cache(maxsize=128)
 def generate_audio(text: str, language_id: str = "en") -> str:
-    payload = {
-        "text": text,
-        "language_id": language_id,
-        "mode": "Speak 🗣️"
-    }
-
+    payload = {"text": text, "language_id": language_id, "mode": "Speak 🗣️"}
     r = SESSION.post(TTS_API_URL, json=payload, timeout=None)
     r.raise_for_status()
 
     wav_path = f"/tmp/tts_{uuid.uuid4().hex}.wav"
 
-    # Case 1: raw audio
+    # raw audio bytes
     if r.headers.get("content-type", "").startswith("audio"):
         with open(wav_path, "wb") as f:
             f.write(r.content)
         return wav_path
 
-    # Case 2: JSON base64
+    # JSON base64
     data = r.json()
     audio_b64 = data.get("audio") or data.get("audio_base64") or data.get("wav")
     if not audio_b64:
@@ -157,21 +158,23 @@ def run_pipeline(question: str, language_id: str):
     if not question.strip():
         return "", None
 
-    answer = answer_question(question)
+    # 1️⃣ Answer text
+    answer = answer_question(question, language_id)
 
+    # 2️⃣ TTS
     try:
         audio_path = generate_audio(answer, language_id)
     except Exception as e:
-        print("TTS generation failed:", e)
+        print("TTS failed:", e)
         audio_path = None
 
     return f"**Bot:** {answer}", audio_path
 
 # =====================================================
-# UI
+# GRADIO UI
 # =====================================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 📄 Qwen CPU Assistant + TTS")
+    gr.Markdown("# 📄 Qwen CPU Assistant + Multilingual TTS")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -197,5 +200,5 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[answer_text, answer_audio]
     )
 
-demo.queue()  # long-running jobs OK (up to 5 min audio)
+demo.queue()  # enable long-running jobs for TTS
 demo.launch(server_name="0.0.0.0", server_port=7860, share=False)