Migrate to faster-whisper with INT8 quantization for ~4x speedup

Dockerfile

@@ -1,7 +1,8 @@
 # ── Tahkik Inference Space ──────────────────────────────────────────────────
-# CPU image. To enable GPU (T4/L4/A100), change the base image to:
+# Uses faster-whisper (CTranslate2 INT8) for ~4x faster inference vs PyTorch.
+# To enable GPU (T4/L4/A100), change the base image to:
 #   FROM nvidia/cuda:12.1-runtime-ubuntu22.04
-# and replace the pip torch line with the CUDA-specific wheel URL.
+# and set compute_type="float16" in main.py.
 # ---------------------------------------------------------------------------
 
 FROM python:3.10-slim
@@ -20,9 +21,8 @@ COPY --chown=user . .
 
 # Redirect all model/cache downloads to /tmp (only writable path in Spaces).
 ENV HF_HOME=/tmp/huggingface_cache
-ENV TORCH_HOME=/tmp/torch_cache
-ENV TRANSFORMERS_VERBOSITY=error
 ENV HF_HUB_DISABLE_PROGRESS_BARS=1
+ENV CT2_VERBOSE=0
 
 USER user
 

main.py

@@ -2,37 +2,34 @@
 """
 Tahkik Inference Server — Hugging Face Space entry point.
 
-Loads the Whisper model ONCE at startup, then serves:
+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 sys
-import struct
 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("TORCH_HOME", "/tmp/torch_cache")
-os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
 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
-import torch
-import torch.nn.functional as F
-from transformers import WhisperForConditionalGeneration, WhisperProcessor
+from faster_whisper import WhisperModel
 
 # ---------------------------------------------------------------------------
 # Constants
 # ---------------------------------------------------------------------------
 
-TAHKIK_MODEL   = "benhadjermed/tahkik-basic-warsh"
+TAHKIK_MODEL   = "benhadjermed/tahkik-small-warsh-ct2"
 SAMPLE_RATE    = 16000
 CHUNK_LENGTH_S = 30
 OVERLAP_S      = 1
@@ -41,39 +38,26 @@ OVERLAP_S      = 1
 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)
+
 ALLOWED_EXTS = {".wav", ".m4a", ".mp3", ".flac", ".ogg"}
 
 # ---------------------------------------------------------------------------
 # Model loading (happens once at module import / server startup)
 # ---------------------------------------------------------------------------
 
-print("[inference] importing torch / transformers...", flush=True)
-device      = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-print(f"[inference] device: {device}", flush=True)
-
-print("[inference] loading processor (openai/whisper-base)...", flush=True)
-processor = WhisperProcessor.from_pretrained(
-    "openai/whisper-base", language="Arabic", task="transcribe"
+print("[inference] loading faster-whisper model...", flush=True)
+model = WhisperModel(
+    TAHKIK_MODEL,
+    device="cpu",
+    compute_type="int8",
+    download_root="/tmp/huggingface_cache",
 )
-
-print(f"[inference] loading model ({TAHKIK_MODEL})...", flush=True)
-model = WhisperForConditionalGeneration.from_pretrained(
-    TAHKIK_MODEL, torch_dtype=torch_dtype
-).to(device)
-
-# Patch missing generation config fields that some fine-tuned checkpoints omit.
-if not hasattr(model.generation_config, "lang_to_id") or model.generation_config.lang_to_id is None:
-    print("[inference] patching generation config from base model...", flush=True)
-    _base = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
-    model.generation_config.lang_to_id = _base.generation_config.lang_to_id
-    model.generation_config.id_to_lang = {v: k for k, v in _base.generation_config.lang_to_id.items()}
-    model.generation_config.task_to_id = _base.generation_config.task_to_id
-    del _base
-
 print("[inference] model ready", flush=True)
 
-# Global inference lock — one inference at a time to avoid GPU OOM.
+# Global inference lock — one inference at a time to avoid resource contention.
 _inference_lock = asyncio.Lock()
 
 # ---------------------------------------------------------------------------
@@ -139,6 +123,7 @@ async def stream_transcribe(ws: WebSocket):
 
     # Accumulate raw PCM bytes from the client.
     audio_buffer = bytearray()
+    session_text = ""
     last_inference_len = 0  # track buffer size at last inference to avoid redundant runs
 
     async def _run_partial(pcm_data: bytes):
@@ -147,9 +132,15 @@ async def stream_transcribe(ws: WebSocket):
                 text = await asyncio.get_event_loop().run_in_executor(
                     None, _transcribe_pcm_buffer, pcm_data
                 )
-            await ws.send_json({"type": "partial", "text": text})
+            full_text = (session_text + " " + text).strip()
+            try:
+                await ws.send_json({"type": "partial", "text": full_text})
+            except Exception:
+                pass  # Connection likely closed
         except Exception as e:
-            print(f"[ws] partial inference error: {e}", flush=True)
+            import traceback
+            err_msg = traceback.format_exc()
+            print(f"[ws] partial inference error:\n{err_msg}", flush=True)
 
     try:
         while True:
@@ -163,20 +154,50 @@ async def stream_transcribe(ws: WebSocket):
                 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 and new_samples >= (SAMPLE_RATE // 2):
-                    # Run partial inference ONLY if the CPU is free.
-                    # This prevents thousands of requests from queuing and timing out the final run.
-                    if not _inference_lock.locked():
-                        last_inference_len = len(audio_buffer)
-                        # We must run this in the background, otherwise we block ws.receive()
-                        asyncio.create_task(_run_partial(bytes(audio_buffer)))
+                if buffer_samples >= MIN_SAMPLES_FOR_INFERENCE:
+                    if _has_trailing_silence(bytes(audio_buffer), SILENCE_THRESHOLD, SILENCE_SAMPLES):
+                        print(f"[ws] auto-finalizing chunk due to silence", flush=True)
+                        async with _inference_lock:
+                            chunk_text = await asyncio.get_event_loop().run_in_executor(
+                                None, _transcribe_pcm_buffer, bytes(audio_buffer)
+                            )
+                        session_text = (session_text + " " + chunk_text).strip()
+                        try:
+                            await ws.send_json({"type": "partial", "text": session_text})
+                        except RuntimeError:
+                            # Client closed connection while we were running inference
+                            break
+
+                        audio_buffer = bytearray()
+                        last_inference_len = 0
+                        continue
+
+                    # Prevent OOM if mic is left on but user is entirely silent for 10s
+                    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 purely silence, dropping...", flush=True)
+                            audio_buffer = bytearray()
+                            last_inference_len = 0
+                            continue
+
+                    if new_samples >= (SAMPLE_RATE // 2):
+                        # Run partial inference ONLY if the lock is free.
+                        # This prevents thousands of requests from queuing and timing out the final run.
+                        if not _inference_lock.locked():
+                            last_inference_len = len(audio_buffer)
+                            # Run in background so ws.receive() is not blocked.
+                            asyncio.create_task(_run_partial(bytes(audio_buffer)))
 
             # --- Text frame: control message ------------------------------
             elif "text" in message and message["text"] is not None:
                 try:
                     msg = json.loads(message["text"])
                 except json.JSONDecodeError:
-                    await ws.send_json({"type": "error", "message": "invalid JSON"})
+                    try:
+                        await ws.send_json({"type": "error", "message": "invalid JSON"})
+                    except RuntimeError:
+                        pass
                     continue
 
                 if msg.get("type") == "stop":
@@ -184,12 +205,15 @@ async def stream_transcribe(ws: WebSocket):
 
                     buffer_samples = len(audio_buffer) // 2
                     if buffer_samples < MIN_SAMPLES_FOR_INFERENCE:
-                        await ws.send_json({
-                            "type": "final",
-                            "text": "",
-                            "confidence": 0.0,
-                            "processing_time_ms": 0,
-                        })
+                        try:
+                            await ws.send_json({
+                                "type": "final",
+                                "text": session_text,
+                                "confidence": 1.0,
+                                "processing_time_ms": 0,
+                            })
+                        except RuntimeError:
+                            pass
                     else:
                         t_start = time.time()
                         async with _inference_lock:
@@ -198,22 +222,28 @@ async def stream_transcribe(ws: WebSocket):
                             )
                         elapsed = int((time.time() - t_start) * 1000)
 
-                        await ws.send_json({
-                            "type": "final",
-                            "text": text,
-                            "confidence": confidence,
-                            "processing_time_ms": elapsed,
-                        })
+                        final_text = (session_text + " " + text).strip()
+                        try:
+                            await ws.send_json({
+                                "type": "final",
+                                "text": final_text,
+                                "confidence": confidence,
+                                "processing_time_ms": elapsed,
+                            })
+                        except RuntimeError:
+                            pass
 
                     # Reset for potential next session on the same connection.
                     audio_buffer = bytearray()
+                    session_text = ""
                     last_inference_len = 0
                     break  # Close after final result.
 
     except WebSocketDisconnect:
         print("[ws] client disconnected", flush=True)
     except Exception as exc:
-        print(f"[ws] error: {exc}", flush=True)
+        import traceback
+        print(f"[ws] error:\n{traceback.format_exc()}", flush=True)
         try:
             await ws.send_json({"type": "error", "message": str(exc)})
         except Exception:
@@ -236,8 +266,32 @@ def _pcm_bytes_to_float32(pcm_bytes: bytes) -> np.ndarray:
     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 _transcribe_pcm_buffer(pcm_bytes: bytes) -> str:
-    """Run Whisper inference on raw PCM buffer, return text only."""
+    """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).
@@ -245,52 +299,41 @@ def _transcribe_pcm_buffer(pcm_bytes: bytes) -> str:
     if len(audio_array) > max_samples:
         audio_array = audio_array[-max_samples:]
 
-    inputs = processor(
-        audio_array, sampling_rate=SAMPLE_RATE, return_tensors="pt"
-    ).input_features.to(device, dtype=torch_dtype)
-
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs,
-            language="ar",
-            task="transcribe",
-        )
-
-    text = processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
-    return text
+    segments, _ = model.transcribe(
+        audio_array,
+        language="ar",
+        task="transcribe",
+        vad_filter=False,
+    )
+    return " ".join(seg.text.strip() for seg in segments)
 
 
 def _transcribe_pcm_buffer_with_confidence(pcm_bytes: bytes) -> tuple:
-    """Run Whisper inference on raw PCM buffer, return (text, confidence)."""
+    """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 = []
 
     for chunk in chunks:
-        inputs = processor(
-            chunk, sampling_rate=SAMPLE_RATE, return_tensors="pt"
-        ).input_features.to(device, dtype=torch_dtype)
-
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs,
-                language="ar",
-                task="transcribe",
-                return_dict_in_generate=True,
-                output_scores=True,
-            )
-
-        text = processor.batch_decode(outputs.sequences, skip_special_tokens=True)[0].strip()
-        all_texts.append(text)
-
-        if outputs.scores:
-            token_probs = [F.softmax(s, dim=-1).max(dim=-1).values for s in outputs.scores]
-            chunk_score = float(sum(p.mean().item() for p in token_probs) / len(token_probs))
+        segments, _ = model.transcribe(
+            chunk,
+            language="ar",
+            task="transcribe",
+            vad_filter=False,
+        )
+        chunk_texts = []
+        chunk_logprobs = []
+        for seg in segments:
+            chunk_texts.append(seg.text.strip())
+            chunk_logprobs.append(seg.avg_logprob)
+
+        all_texts.append(" ".join(chunk_texts))
+        if chunk_logprobs:
+            avg = sum(chunk_logprobs) / len(chunk_logprobs)
+            all_scores.append(_logprob_to_confidence(avg))
         else:
-            chunk_score = 1.0
-        all_scores.append(chunk_score)
+            all_scores.append(1.0)
 
     transcription = " ".join(all_texts)
     confidence = round(sum(all_scores) / len(all_scores), 4) if all_scores else 0.0
@@ -316,39 +359,35 @@ def _split_audio(audio_array, sr=SAMPLE_RATE, chunk_s=CHUNK_LENGTH_S, overlap_s=
 def _transcribe_file(audio_path: str) -> dict:
     import librosa
 
-    t_start     = time.time()
+    t_start = time.time()
     audio_array, _ = librosa.load(audio_path, sr=SAMPLE_RATE)
 
-    chunks     = _split_audio(audio_array)
-    all_texts  = []
+    chunks = _split_audio(audio_array)
+    all_texts = []
     all_scores = []
 
     for chunk in chunks:
-        inputs = processor(
-            chunk, sampling_rate=SAMPLE_RATE, return_tensors="pt"
-        ).input_features.to(device, dtype=torch_dtype)
-
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs,
-                language="ar",
-                task="transcribe",
-                return_dict_in_generate=True,
-                output_scores=True,
-            )
-
-        text = processor.batch_decode(outputs.sequences, skip_special_tokens=True)[0].strip()
-        all_texts.append(text)
-
-        if outputs.scores:
-            token_probs = [F.softmax(s, dim=-1).max(dim=-1).values for s in outputs.scores]
-            chunk_score = float(sum(p.mean().item() for p in token_probs) / len(token_probs))
+        segments, _ = model.transcribe(
+            chunk,
+            language="ar",
+            task="transcribe",
+            vad_filter=False,
+        )
+        chunk_texts = []
+        chunk_logprobs = []
+        for seg in segments:
+            chunk_texts.append(seg.text.strip())
+            chunk_logprobs.append(seg.avg_logprob)
+
+        all_texts.append(" ".join(chunk_texts))
+        if chunk_logprobs:
+            avg = sum(chunk_logprobs) / len(chunk_logprobs)
+            all_scores.append(_logprob_to_confidence(avg))
         else:
-            chunk_score = 1.0
-        all_scores.append(chunk_score)
+            all_scores.append(1.0)
 
     return {
-        "transcription":     " ".join(all_texts),
-        "confidence_score":  round(sum(all_scores) / len(all_scores), 4) if all_scores else 0.0,
+        "transcription":      " ".join(all_texts),
+        "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),
     }

requirements.txt

@@ -1,9 +1,7 @@
 fastapi
 uvicorn[standard]
-torch
-transformers
+faster-whisper
 librosa
 soundfile
-accelerate
 python-multipart
 numpy