#!/usr/bin/env python3 """ Tahkik Inference Server — Hugging Face Space entry point. Loads the Whisper model ONCE at startup via faster-whisper (CTranslate2), then serves: - POST /evaluate — batch transcription (upload a full audio file) - WS /ws/stream — real-time streaming transcription (send PCM chunks) """ import asyncio import json import math import os import time import tempfile # Redirect model caches to /tmp (only writable dir in HF Spaces) os.environ.setdefault("HF_HOME", "/tmp/huggingface_cache") os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") os.environ.setdefault("CT2_VERBOSE", "0") import numpy as np from fastapi import FastAPI, File, UploadFile, HTTPException, WebSocket, WebSocketDisconnect from fastapi.responses import JSONResponse from faster_whisper import WhisperModel # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- TAHKIK_MODEL = "benhadjermed/tahkik-basic-warsh" SAMPLE_RATE = 16000 CHUNK_LENGTH_S = 30 OVERLAP_S = 1 # Minimum seconds of audio before running partial inference (reduces hallucinations). # Whisper produces unreliable text on snippets shorter than ~1 s. MIN_AUDIO_FOR_INFERENCE_S = 1.0 MIN_SAMPLES_FOR_INFERENCE = int(MIN_AUDIO_FOR_INFERENCE_S * SAMPLE_RATE) SILENCE_THRESHOLD = 0.02 # RMS threshold for silence SILENCE_DURATION_S = 0.8 # seconds of trailing silence to trigger finalization SILENCE_SAMPLES = int(SILENCE_DURATION_S * SAMPLE_RATE) # Max characters of context (initial_prompt) we accept from the client. # Long prompts cause Whisper to hallucinate the prompt back on silence — # 200 chars is enough to bias vocabulary toward the current ayah's words # without overwhelming the audio signal. MAX_CONTEXT_CHARS = 200 # faster-whisper transcribe options shared by every inference call. # These are the standard anti-hallucination knobs — see openai/whisper#679, # guillaumekln/faster-whisper#108 for background. WHISPER_OPTS = dict( language="ar", task="transcribe", # Lightweight VAD — keeps word endings (the previous concern that led # to vad_filter=False) but strips long silence chunks the model would # otherwise hallucinate into. vad_filter=True, vad_parameters=dict(min_silence_duration_ms=300, threshold=0.35), # Standard temperature-fallback chain. When a decode fails the # compression-ratio or log-prob check below, faster-whisper retries # at the next temperature. After exhausting all values it drops the # segment entirely. temperature=[0.0, 0.2, 0.4, 0.6, 0.8, 1.0], compression_ratio_threshold=2.4, # detects loop hallucinations log_prob_threshold=-1.0, # detects low-quality decodes no_speech_threshold=0.6, # treat segment as silence above this # Each partial is a fresh window — don't carry decoder state forward. # This eliminates loop hallucinations like وَعِيسيا وَعِيسيا وَعِيسيا. condition_on_previous_text=False, ) # Drop any segment the model itself flagged as likely non-speech. NO_SPEECH_PROB_DROP_THRESHOLD = 0.7 # Common Whisper hallucinations on silence / non-speech audio. faster-whisper # trained on YouTube emits these reliably when given silence; we drop them. HALLUCINATION_PHRASES = ( "شكرا لك", "شكراً لك", "شكرا لكم", "شكراً لكم", "شكرا لمشاهدتكم", "شكراً لمشاهدتكم", "ترجمة نانسي قنقر", "ترجمة", "تابعونا", "اشتراك في القناة", "Thank you", "thank you", ".", ) ALLOWED_EXTS = {".wav", ".m4a", ".mp3", ".flac", ".ogg"} # --------------------------------------------------------------------------- # Model loading (happens once at module import / server startup) # --------------------------------------------------------------------------- CT2_CONVERTED_DIR = f"/tmp/tahkik-ct2-{TAHKIK_MODEL.replace('/', '_')}" def _ensure_ct2_model(hf_model_id: str, ct2_dir: str) -> str: """ Return a CTranslate2 model directory for the given HF model ID. If the model is already in CT2 format (has model.bin), return ct2_dir directly after downloading via faster-whisper's normal path. Otherwise treat it as a Transformers Whisper checkpoint, convert it to CT2 int8 once, cache the result in ct2_dir, and return that path. """ import shutil from huggingface_hub import snapshot_download # Try downloading as a CT2 model first (works if repo already has model.bin) hf_cache = os.environ.get("HF_HOME", "/tmp/huggingface_cache") local_hf = snapshot_download(repo_id=hf_model_id, cache_dir=hf_cache) if os.path.exists(os.path.join(local_hf, "model.bin")): print(f"[inference] model is already CTranslate2 format, loading directly", flush=True) return local_hf # PyTorch / safetensors checkpoint — convert to CT2 int8. model_bin = os.path.join(ct2_dir, "model.bin") if not os.path.exists(model_bin): print(f"[inference] converting {hf_model_id} → CTranslate2 int8 in {ct2_dir}...", flush=True) from ctranslate2.converters import TransformersConverter converter = TransformersConverter( model_name_or_path=local_hf, low_cpu_mem_usage=True, ) os.makedirs(ct2_dir, exist_ok=True) converter.convert(ct2_dir, quantization="int8", force=True) print("[inference] conversion complete", flush=True) else: print(f"[inference] using cached CT2 model at {ct2_dir}", flush=True) return ct2_dir print("[inference] loading faster-whisper model...", flush=True) _ct2_path = _ensure_ct2_model(TAHKIK_MODEL, CT2_CONVERTED_DIR) model = WhisperModel( _ct2_path, device="cpu", compute_type="int8", ) print("[inference] model ready", flush=True) # Global inference lock — one inference at a time to avoid resource contention. _inference_lock = asyncio.Lock() # --------------------------------------------------------------------------- # FastAPI app # --------------------------------------------------------------------------- app = FastAPI(title="Tahkik Inference API") @app.get("/health") def health(): return {"status": "ok"} # --------------------------------------------------------------------------- # POST /evaluate — batch transcription (backward compatible) # --------------------------------------------------------------------------- @app.post("/evaluate") async def evaluate(audio: UploadFile = File(...)): filename = audio.filename or "recording.wav" ext = os.path.splitext(filename)[1].lower() or ".wav" if ext not in ALLOWED_EXTS: raise HTTPException(status_code=400, detail=f"unsupported audio format: {ext}") data = await audio.read() with tempfile.NamedTemporaryFile(suffix=ext, delete=False, dir="/tmp") as f: f.write(data) tmp_path = f.name try: result = _transcribe_file(tmp_path) except Exception as exc: raise HTTPException(status_code=500, detail=str(exc)) finally: os.unlink(tmp_path) return JSONResponse(result) # --------------------------------------------------------------------------- # WS /ws/stream — real-time streaming transcription # --------------------------------------------------------------------------- @app.websocket("/ws/stream") async def stream_transcribe(ws: WebSocket): """ Real-time streaming transcription over WebSocket. Protocol: Client → Server: - Binary frames: raw PCM 16-bit signed LE, 16 kHz, mono - Text frame: JSON {"type": "stop"} to signal end of recording - Text frame: JSON {"type": "ping"} keepalive heartbeat - Text frame: JSON {"type": "context", "text": "..."} — canonical ayah text the user is reciting (used as initial_prompt to bias decoding and reduce hallucinations). Server → Client: - Text frames: JSON messages {"type": "partial", "text": "..."} — intermediate transcription {"type": "final", "text": "...", "confidence": 0.94, "processing_time_ms": 1234} {"type": "error", "message": "..."} {"type": "pong"} — heartbeat response """ await ws.accept() print("[ws] client connected", flush=True) # Track background inference tasks so we can cancel them on disconnect. _pending_tasks: list[asyncio.Task] = [] # Generation counter — incremented on every state reset. Background tasks # check this before writing results so stale tasks discard their output. _session_gen = 0 # Accumulate raw PCM bytes from the client. audio_buffer = bytearray() session_text = "" _last_sent_text = "" # Deduplication: only send partial if text changed last_inference_len = 0 # track buffer size at last inference to avoid redundant runs # ── Incremental inference tracking ── # Only transcribe audio AFTER this byte offset for partials. # This keeps partial inference time constant regardless of total buffer size. last_partial_offset = 0 # ── First-partial optimization ── # Use a lower threshold for the very first partial to show text faster. first_partial_done = False # ── Decoding bias ── # Set by the client via {"type": "context", ...} once it has detected # which ayah the user is reciting. Whisper uses initial_prompt to bias # the decoder toward expected text — same word for same syllables. initial_prompt: str | None = None # Sliding window size for partial transcription. Each partial transcribes # the LAST N seconds of audio as a single contiguous window and the result # REPLACES session_text. This avoids the chunk-boundary word-cutting bug # where transcribing 0.5s chunks separately produces broken words at every # boundary (the start of one chunk = mid-word of the previous chunk). PARTIAL_WINDOW_S = 8 async def _run_partial_window(pcm_data: bytes, gen: int): """Transcribe the last N seconds of audio as one contiguous window.""" nonlocal session_text, _last_sent_text try: if len(pcm_data) < MIN_SAMPLES_FOR_INFERENCE * 2: return # Take the trailing window — words at boundaries are complete. window_bytes = SAMPLE_RATE * PARTIAL_WINDOW_S * 2 audio = pcm_data[-window_bytes:] if len(pcm_data) > window_bytes else pcm_data if _is_silent(audio): return async with _inference_lock: text = await asyncio.get_event_loop().run_in_executor( None, _transcribe_pcm_buffer, audio, initial_prompt ) # Discard result if the session was reset during inference. if gen != _session_gen: return text = _filter_hallucinations(text).strip() if not text: return # REPLACE session_text — the window transcription is the full text # for that window, no concatenation needed. session_text = text if session_text != _last_sent_text: _last_sent_text = session_text print(f"[ws] partial: {session_text[:80]}...", flush=True) try: await ws.send_json({"type": "partial", "text": session_text}) except Exception: pass except asyncio.CancelledError: pass except Exception: import traceback print(f"[ws] partial inference error:\n{traceback.format_exc()}", flush=True) try: while True: try: message = await ws.receive() except (WebSocketDisconnect, RuntimeError): # Client closed (RuntimeError = "Cannot call 'receive' once a # disconnect message has been received"). Just exit cleanly. break # --- Binary frame: audio chunk -------------------------------- if "bytes" in message and message["bytes"] is not None: audio_buffer.extend(message["bytes"]) # Only run inference if we have enough new audio. buffer_samples = len(audio_buffer) // 2 # 16-bit = 2 bytes/sample new_samples = buffer_samples - (last_inference_len // 2) if buffer_samples >= MIN_SAMPLES_FOR_INFERENCE: # Prevent OOM if mic is left open with pure silence for 10s. # Only drop the audio buffer — session_text is preserved. if buffer_samples > SAMPLE_RATE * 10: audio_array = _pcm_bytes_to_float32(bytes(audio_buffer)) if np.sqrt(np.mean(audio_array ** 2)) < SILENCE_THRESHOLD * 2: print("[ws] buffer full of pure silence, dropping audio...", flush=True) audio_buffer = bytearray() last_inference_len = 0 last_partial_offset = 0 continue # Trigger partial as fast as inference can keep up. # Need at least 0.3s of new audio to avoid wasted inferences. min_new_samples = SAMPLE_RATE // 3 if new_samples >= min_new_samples and not _inference_lock.locked(): first_partial_done = True last_inference_len = len(audio_buffer) task = asyncio.create_task( _run_partial_window(bytes(audio_buffer), _session_gen) ) _pending_tasks.append(task) task.add_done_callback(lambda t: _pending_tasks.remove(t) if t in _pending_tasks else None) # --- Text frame: control message ------------------------------ elif "text" in message and message["text"] is not None: try: msg = json.loads(message["text"]) except json.JSONDecodeError: try: await ws.send_json({"type": "error", "message": "invalid JSON"}) except RuntimeError: pass continue # ── Ping/pong heartbeat ── if msg.get("type") == "ping": try: await ws.send_json({"type": "pong"}) except RuntimeError: pass continue # ── Session reset ── if msg.get("type") == "reset": for task in list(_pending_tasks): task.cancel() if _pending_tasks: await asyncio.gather(*_pending_tasks, return_exceptions=True) _session_gen += 1 audio_buffer = bytearray() session_text = "" _last_sent_text = "" last_inference_len = 0 last_partial_offset = 0 first_partial_done = False print("[ws] session reset by client", flush=True) continue # ── Decoding bias ── if msg.get("type") == "context": ctx_text = (msg.get("text") or "").strip() if ctx_text: initial_prompt = ctx_text[:MAX_CONTEXT_CHARS] print(f"[ws] context set ({len(initial_prompt)} chars)", flush=True) else: initial_prompt = None continue if msg.get("type") == "stop": print(f"[ws] stop received, buffer size: {len(audio_buffer)} bytes", flush=True) # Cancel background tasks before final inference — they must # not update session_text after we start the final pass. for task in list(_pending_tasks): task.cancel() if _pending_tasks: await asyncio.gather(*_pending_tasks, return_exceptions=True) _session_gen += 1 # Send the partial-accumulated result immediately — client # already uses this; the server "final" is only refinement. try: await ws.send_json({ "type": "final", "text": session_text, "confidence": 1.0, "processing_time_ms": 0, }) except RuntimeError: pass # Refine in background using the full audio buffer (capped # at 30s — Whisper's context limit). The result REPLACES # session_text since it's a single contiguous transcription. full_audio = bytes(audio_buffer) if len(full_audio) // 2 >= MIN_SAMPLES_FOR_INFERENCE: max_full_bytes = SAMPLE_RATE * 30 * 2 if len(full_audio) > max_full_bytes: full_audio = full_audio[-max_full_bytes:] _refine_gen = _session_gen # capture before reset _refine_prompt = initial_prompt _refine_partial = session_text async def _refine_final(audio=full_audio, prompt=_refine_prompt, gen=_refine_gen, partial=_refine_partial): async with _inference_lock: t, conf = await asyncio.get_event_loop().run_in_executor( None, _transcribe_pcm_buffer_with_confidence, audio, prompt ) if gen != _session_gen: return refined = _filter_hallucinations(t).strip() if refined and refined != partial: try: await ws.send_json({ "type": "final", "text": refined, "confidence": conf, "processing_time_ms": 0, }) except Exception: pass asyncio.create_task(_refine_final()) # Reset state — connection stays open for client to close. audio_buffer = bytearray() session_text = "" _last_sent_text = "" last_inference_len = 0 last_partial_offset = 0 first_partial_done = False initial_prompt = None # Don't break — let the client close the connection. # This avoids a race where the final message hasn't been # received by the client before the server closes the socket. except WebSocketDisconnect: print("[ws] client disconnected", flush=True) except Exception as exc: import traceback print(f"[ws] error:\n{traceback.format_exc()}", flush=True) try: await ws.send_json({"type": "error", "message": str(exc)}) except Exception: pass finally: # Cancel any background inference tasks still running for this connection. # Without this, tasks hold _inference_lock and starve the next connection. for task in list(_pending_tasks): task.cancel() if _pending_tasks: await asyncio.gather(*_pending_tasks, return_exceptions=True) try: await ws.close() except Exception: pass print("[ws] connection closed", flush=True) # --------------------------------------------------------------------------- # Inference helpers # --------------------------------------------------------------------------- def _pcm_bytes_to_float32(pcm_bytes: bytes) -> np.ndarray: """Convert raw PCM 16-bit signed LE bytes to float32 numpy array in [-1, 1].""" int16_array = np.frombuffer(pcm_bytes, dtype=np.int16) return int16_array.astype(np.float32) / 32768.0 def _has_trailing_silence(pcm_bytes: bytes, threshold: float, duration_samples: int) -> bool: """Check if buffer ends with N seconds of silence below threshold, AND had speech before it.""" if len(pcm_bytes) < duration_samples * 2: return False audio_array = _pcm_bytes_to_float32(pcm_bytes) trailing = audio_array[-duration_samples:] rms = np.sqrt(np.mean(trailing ** 2)) if rms < threshold: # Require some actual speech before the trailing silence to count as "trailing silence" leading = audio_array[:-duration_samples] if len(leading) > 0: leading_rms = np.sqrt(np.mean(leading ** 2)) if leading_rms > threshold * 1.5: return True return False def _logprob_to_confidence(avg_logprob: float) -> float: """Convert faster-whisper's avg_logprob to a 0-1 confidence score via exp().""" return math.exp(max(avg_logprob, -5.0)) # clamp to avoid exp(-inf) = 0 def _is_silent(pcm_bytes: bytes, threshold: float = SILENCE_THRESHOLD) -> bool: """Return True if audio RMS is below threshold (effectively silence).""" if not pcm_bytes: return True audio = _pcm_bytes_to_float32(pcm_bytes) return float(np.sqrt(np.mean(audio ** 2))) < threshold def _filter_hallucinations(text: str) -> str: """Strip known Whisper hallucination phrases from transcription output.""" if not text: return text result = text.strip() for phrase in HALLUCINATION_PHRASES: result = result.replace(phrase, "").strip() return result def _transcribe_pcm_buffer(pcm_bytes: bytes, initial_prompt: str | None = None) -> str: """Run faster-whisper inference on raw PCM buffer, return text only.""" audio_array = _pcm_bytes_to_float32(pcm_bytes) # Limit to last 30 seconds (Whisper's context window). max_samples = CHUNK_LENGTH_S * SAMPLE_RATE if len(audio_array) > max_samples: audio_array = audio_array[-max_samples:] prompt = initial_prompt[-MAX_CONTEXT_CHARS:] if initial_prompt else None segments, _ = model.transcribe( audio_array, initial_prompt=prompt, **WHISPER_OPTS, ) parts = [ seg.text.strip() for seg in segments if seg.no_speech_prob < NO_SPEECH_PROB_DROP_THRESHOLD ] return " ".join(p for p in parts if p) def _transcribe_pcm_buffer_with_confidence(pcm_bytes: bytes, initial_prompt: str | None = None) -> tuple: """Run faster-whisper inference on raw PCM buffer, return (text, confidence).""" audio_array = _pcm_bytes_to_float32(pcm_bytes) chunks = _split_audio(audio_array) all_texts = [] all_scores = [] prompt = initial_prompt[-MAX_CONTEXT_CHARS:] if initial_prompt else None for chunk in chunks: segments, _ = model.transcribe( chunk, initial_prompt=prompt, **WHISPER_OPTS, ) chunk_texts = [] chunk_logprobs = [] for seg in segments: if seg.no_speech_prob >= NO_SPEECH_PROB_DROP_THRESHOLD: continue chunk_texts.append(seg.text.strip()) chunk_logprobs.append(seg.avg_logprob) all_texts.append(" ".join(t for t in chunk_texts if t)) if chunk_logprobs: avg = sum(chunk_logprobs) / len(chunk_logprobs) all_scores.append(_logprob_to_confidence(avg)) else: all_scores.append(1.0) transcription = " ".join(t for t in all_texts if t) confidence = round(sum(all_scores) / len(all_scores), 4) if all_scores else 0.0 return transcription, confidence def _split_audio(audio_array, sr=SAMPLE_RATE, chunk_s=CHUNK_LENGTH_S, overlap_s=OVERLAP_S): chunk_len = int(chunk_s * sr) step_len = int((chunk_s - overlap_s) * sr) chunks = [] start = 0 while start < len(audio_array): end = min(start + chunk_len, len(audio_array)) chunks.append(audio_array[start:end]) start += step_len remaining = len(audio_array) - start if 0 < remaining < 2 * sr: chunks[-1] = audio_array[start - step_len:] break return chunks def _transcribe_file(audio_path: str) -> dict: import librosa t_start = time.time() audio_array, _ = librosa.load(audio_path, sr=SAMPLE_RATE) chunks = _split_audio(audio_array) all_texts = [] all_scores = [] for chunk in chunks: segments, _ = model.transcribe(chunk, **WHISPER_OPTS) chunk_texts = [] chunk_logprobs = [] for seg in segments: if seg.no_speech_prob >= NO_SPEECH_PROB_DROP_THRESHOLD: continue chunk_texts.append(seg.text.strip()) chunk_logprobs.append(seg.avg_logprob) all_texts.append(" ".join(t for t in chunk_texts if t)) if chunk_logprobs: avg = sum(chunk_logprobs) / len(chunk_logprobs) all_scores.append(_logprob_to_confidence(avg)) else: all_scores.append(1.0) return { "transcription": " ".join(t for t in all_texts if t), "confidence_score": round(sum(all_scores) / len(all_scores), 4) if all_scores else 0.0, "processing_time_ms": int((time.time() - t_start) * 1000), }