diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..6bfe385f7624d1638d3fc9535f43fcc66ca491d2 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,16 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/1bf8e725-d7c9-40b1-ba8f-f87711c6f146/part_0_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/4036b9ca-15fe-43b1-ada3-eb7603ddda14/part_0_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/8c701121-c985-485b-aaac-a1ff6805ad15/part_0_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/bc5cf2ff-eef5-41f3-8c32-e623b0771590/part_0_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/d0ae91ba-acf5-44fb-8863-fc9efad5ef09/part_1_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/f9943ee7-cbfc-4972-9b4f-c5737473b88a/part_0_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/_indices/fe5c2cf5-6aad-44b5-8308-a08422c6a259/part_0_invert.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/data/000101001101011011001000d12f72422994e6d4750671d1cb.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/data/000111110111000011000010fb6b6343448c22c242e8a7e820.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/data/1001010110001010001010117460c543a7ab8135cc05eb0a67.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/data/100110110110110110000101ab4f8a403c88cf44f5dfcc16e2.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/data/1110110000011010101111109886fb45d299593d03bffd0230.lance filter=lfs diff=lfs merge=lfs -text +lancedb/passages.lance/data/1111001101101101110100007dd4244f748f751fd19c94e429.lance filter=lfs diff=lfs merge=lfs -text diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..58f94b998286144dac1dddcab1d3e86f4fd46bbf --- /dev/null +++ b/Dockerfile @@ -0,0 +1,23 @@ +# Slim online SearchAudio image (CPU only). Build context = project root: +# docker build -f deploy/Dockerfile -t searchaudio-online . +FROM python:3.11-slim + +WORKDIR /app +ENV HF_HOME=/app/models \ + HF_HUB_CACHE=/app/models/hub \ + SEARCHAUDIO_AUDIO_MODE=archive \ + PYTHONUNBUFFERED=1 + +COPY requirements-online.txt . +RUN pip install --no-cache-dir -r requirements-online.txt + +# Bake bge-m3 into the image so cold starts don't download it; then run offline. +RUN python -c "from FlagEmbedding import BGEM3FlagModel; BGEM3FlagModel('BAAI/bge-m3', use_fp16=False)" +ENV HF_HUB_OFFLINE=1 + +COPY app ./app +COPY config.yaml ./config.yaml +COPY deploy/index/lancedb ./data/index/lancedb + +EXPOSE 7860 +CMD ["uvicorn", "app.server:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"] diff --git a/app/__init__.py b/app/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..133d726689835e3a72044ddd6db914d4f8d10ad0 --- /dev/null +++ b/app/__init__.py @@ -0,0 +1,12 @@ +"""SearchAudio — fully-local semantic search over long Hindi audio talks. + +Type an English sentence (or speak/upload a Hindi clip) and jump straight to the +moment in a recording where that topic is discussed. Everything runs offline on a +single Windows PC with an NVIDIA GPU. + +Pipeline: audio -> ffmpeg normalize -> WhisperX ASR (word timestamps) -> + sentence-packed chunks -> bge-m3 embeddings -> LanceDB index. +Query: text/Hindi-audio -> bge-m3 -> LanceDB hybrid search -> bge-reranker -> results. +""" + +__version__ = "0.1.0" diff --git a/app/__pycache__/__init__.cpython-312.pyc b/app/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..295269e1911ddbb974d476105b95766143920bf9 Binary files /dev/null and b/app/__pycache__/__init__.cpython-312.pyc differ diff --git a/app/__pycache__/__init__.cpython-314.pyc b/app/__pycache__/__init__.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4822c8f215a1eb53357a7e353bdd8ae8932cbee4 Binary files /dev/null and b/app/__pycache__/__init__.cpython-314.pyc differ diff --git a/app/__pycache__/_cuda_bootstrap.cpython-312.pyc b/app/__pycache__/_cuda_bootstrap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6889b792b670d39c6af9447f5a47e4722daee876 Binary files /dev/null and b/app/__pycache__/_cuda_bootstrap.cpython-312.pyc differ diff --git a/app/__pycache__/_cuda_bootstrap.cpython-314.pyc b/app/__pycache__/_cuda_bootstrap.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b45f557ffb2139cd455456fda386528683640036 Binary files /dev/null and b/app/__pycache__/_cuda_bootstrap.cpython-314.pyc differ diff --git a/app/__pycache__/asr.cpython-312.pyc b/app/__pycache__/asr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c0bd74141bc755a4060a9460da6a6b29375c367 Binary files /dev/null and b/app/__pycache__/asr.cpython-312.pyc differ diff --git a/app/__pycache__/asr.cpython-314.pyc b/app/__pycache__/asr.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24bef61df9fe62bbb1d339093000f1198a9c5b6b Binary files /dev/null and b/app/__pycache__/asr.cpython-314.pyc differ diff --git a/app/__pycache__/audio_source.cpython-312.pyc b/app/__pycache__/audio_source.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edba36d12792db96a839933e24407cedcee75410 Binary files /dev/null and b/app/__pycache__/audio_source.cpython-312.pyc differ diff --git a/app/__pycache__/audio_utils.cpython-312.pyc b/app/__pycache__/audio_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60a68ec85f4144243f6012e15037272d167a7dc3 Binary files /dev/null and b/app/__pycache__/audio_utils.cpython-312.pyc differ diff --git a/app/__pycache__/audio_utils.cpython-314.pyc b/app/__pycache__/audio_utils.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27079d45163ad53d84d37f5483a83b92a546bf67 Binary files /dev/null and b/app/__pycache__/audio_utils.cpython-314.pyc differ diff --git a/app/__pycache__/chunking.cpython-312.pyc b/app/__pycache__/chunking.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..734985ccbbe6733c544a26e48dc8c7a3749f20c6 Binary files /dev/null and b/app/__pycache__/chunking.cpython-312.pyc differ diff --git a/app/__pycache__/chunking.cpython-314.pyc b/app/__pycache__/chunking.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3fecf42553423e598d88e925e897729291f97a4a Binary files /dev/null and b/app/__pycache__/chunking.cpython-314.pyc differ diff --git a/app/__pycache__/config.cpython-312.pyc b/app/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c7ab146badb1831579caf58c0926135583c5669 Binary files /dev/null and b/app/__pycache__/config.cpython-312.pyc differ diff --git a/app/__pycache__/config.cpython-314.pyc b/app/__pycache__/config.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fbf870bbcdc4768dd5402c8da85b480966572e20 Binary files /dev/null and b/app/__pycache__/config.cpython-314.pyc differ diff --git a/app/__pycache__/embed.cpython-312.pyc b/app/__pycache__/embed.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..90d0864cb534b9b947f1fa64217c401b3bde4d3f Binary files /dev/null and b/app/__pycache__/embed.cpython-312.pyc differ diff --git a/app/__pycache__/embed.cpython-314.pyc b/app/__pycache__/embed.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..015526065ddbccd68e9c2b512000385336dfec21 Binary files /dev/null and b/app/__pycache__/embed.cpython-314.pyc differ diff --git a/app/__pycache__/ingest.cpython-312.pyc b/app/__pycache__/ingest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92b5b393c71be8f807a2399f00fed60c5c3b8035 Binary files /dev/null and b/app/__pycache__/ingest.cpython-312.pyc differ diff --git a/app/__pycache__/ingest.cpython-314.pyc b/app/__pycache__/ingest.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f95ce77d599cf5e8b0a96029e7ef037ab1e58ec Binary files /dev/null and b/app/__pycache__/ingest.cpython-314.pyc differ diff --git a/app/__pycache__/manifest.cpython-312.pyc b/app/__pycache__/manifest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..91d91c29de71e0ea4bff3b7f2bdd41e3b90f6203 Binary files /dev/null and b/app/__pycache__/manifest.cpython-312.pyc differ diff --git a/app/__pycache__/manifest.cpython-314.pyc b/app/__pycache__/manifest.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f7294560d8dee3575a9daddcbbb524b62e80993 Binary files /dev/null and b/app/__pycache__/manifest.cpython-314.pyc differ diff --git a/app/__pycache__/models.cpython-312.pyc b/app/__pycache__/models.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56669394126aac1e45bdb4b07e7b63278d6c81c9 Binary files /dev/null and b/app/__pycache__/models.cpython-312.pyc differ diff --git a/app/__pycache__/models.cpython-314.pyc b/app/__pycache__/models.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5515d5baa886d1a330c5f8ec8cb185743d86c99d Binary files /dev/null and b/app/__pycache__/models.cpython-314.pyc differ diff --git a/app/__pycache__/query.cpython-312.pyc b/app/__pycache__/query.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..532ab5c39c17df5cc6a66edf708dbd6f4bb180b9 Binary files /dev/null and b/app/__pycache__/query.cpython-312.pyc differ diff --git a/app/__pycache__/query.cpython-314.pyc b/app/__pycache__/query.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4e9a5a18a2b567f43e9018fade4858ec30808064 Binary files /dev/null and b/app/__pycache__/query.cpython-314.pyc differ diff --git a/app/__pycache__/rerank.cpython-312.pyc b/app/__pycache__/rerank.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..313f6ea2d681d25f3fcddf65b69e0b299d5a160b Binary files /dev/null and b/app/__pycache__/rerank.cpython-312.pyc differ diff --git a/app/__pycache__/rerank.cpython-314.pyc b/app/__pycache__/rerank.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d35a8c9f2a08f284bd9e14e14fd1e7d4a26a317 Binary files /dev/null and b/app/__pycache__/rerank.cpython-314.pyc differ diff --git a/app/__pycache__/server.cpython-312.pyc b/app/__pycache__/server.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19fb970e763e0ba6f9b73cdfa42b77f9e35f4548 Binary files /dev/null and b/app/__pycache__/server.cpython-312.pyc differ diff --git a/app/__pycache__/server.cpython-314.pyc b/app/__pycache__/server.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38a204013641ca985275ec208f109a05faee1fb6 Binary files /dev/null and b/app/__pycache__/server.cpython-314.pyc differ diff --git a/app/__pycache__/store.cpython-312.pyc b/app/__pycache__/store.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..510acf0198784264669cd26247cd8f75a6ea7479 Binary files /dev/null and b/app/__pycache__/store.cpython-312.pyc differ diff --git a/app/__pycache__/store.cpython-314.pyc b/app/__pycache__/store.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c66bd74dbf3cb426b09b5b8cc5f181d22de7fd9 Binary files /dev/null and b/app/__pycache__/store.cpython-314.pyc differ diff --git a/app/__pycache__/translate.cpython-314.pyc b/app/__pycache__/translate.cpython-314.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f29bcc15aa084042663f11c83c6922c5a0d4025c Binary files /dev/null and b/app/__pycache__/translate.cpython-314.pyc differ diff --git a/app/_cuda_bootstrap.py b/app/_cuda_bootstrap.py new file mode 100644 index 0000000000000000000000000000000000000000..537461c4984d99770b59db79c80166c52577d5fd --- /dev/null +++ b/app/_cuda_bootstrap.py @@ -0,0 +1,88 @@ +"""Make Windows able to find the CUDA / cuDNN DLLs that CTranslate2 needs. + +The single most common failure when running faster-whisper / WhisperX on Windows is: + + RuntimeError: Library cudnn_ops64_9.dll is not found ... + (or: Could not locate cudnn_ops_infer64_8.dll) + +It happens because CTranslate2 loads cuDNN/cuBLAS by name via the OS loader, but the +DLLs live inside pip wheels (``nvidia-cudnn-cu12`` / ``nvidia-cublas-cu12``) or inside +``torch/lib`` (the cu128 wheel bundles them) — directories that are NOT on the default +DLL search path. + +Importing this module **first** — before importing ``ctranslate2``, ``faster_whisper`` +or ``whisperx`` — registers those directories with the OS loader using +``os.add_dll_directory`` (and prepends them to PATH as a belt-and-suspenders fallback). + +Usage:: + + from app import _cuda_bootstrap # noqa: F401 (must come first) + from faster_whisper import WhisperModel # now finds cuDNN + +This is a no-op on non-Windows platforms. +""" +from __future__ import annotations + +import importlib.util +import os +import sys + +# cuDNN must come last so its directory is searched first (add_dll_directory is LIFO-ish +# in practice via PATH ordering); the order here is the order we *append* candidates. +_NVIDIA_PKGS = ( + "nvidia.cuda_runtime", + "nvidia.cuda_nvrtc", + "nvidia.cublas", + "nvidia.cudnn", +) + + +def _add_dll_dir(path: str) -> bool: + if not path or not os.path.isdir(path): + return False + try: + os.add_dll_directory(path) # type: ignore[attr-defined] (Windows only) + except (OSError, AttributeError): + return False + # Some libraries resolve via LoadLibrary(name) which consults PATH, not the + # add_dll_directory list — so prepend to PATH too. + if path not in os.environ.get("PATH", ""): + os.environ["PATH"] = path + os.pathsep + os.environ.get("PATH", "") + return True + + +def _candidate_dirs() -> list[str]: + dirs: list[str] = [] + for pkg in _NVIDIA_PKGS: + try: + spec = importlib.util.find_spec(pkg) + except (ImportError, ValueError, ModuleNotFoundError): + spec = None + if spec and spec.submodule_search_locations: + base = list(spec.submodule_search_locations)[0] + # Windows wheels put DLLs in /bin, Linux in /lib — add both, harmless if absent. + dirs.append(os.path.join(base, "bin")) + dirs.append(os.path.join(base, "lib")) + # torch's bundled CUDA libs (the cu128 wheel ships cudnn/cublas DLLs here). + try: + import torch # noqa: WPS433 (local import on purpose) + + dirs.append(os.path.join(os.path.dirname(torch.__file__), "lib")) + except Exception: # noqa: BLE001 — torch may not be importable yet; ignore. + pass + return dirs + + +def setup_cuda_dll_path() -> list[str]: + """Register CUDA/cuDNN DLL directories with the Windows loader. Returns dirs added.""" + if sys.platform != "win32": + return [] + added: list[str] = [] + for d in _candidate_dirs(): + if _add_dll_dir(d): + added.append(d) + return added + + +# Run on import. Safe to import multiple times. +ADDED_DLL_DIRS = setup_cuda_dll_path() diff --git a/app/asr.py b/app/asr.py new file mode 100644 index 0000000000000000000000000000000000000000..ae55010366095a0e05b5e90571b2bf7354920049 --- /dev/null +++ b/app/asr.py @@ -0,0 +1,209 @@ +"""Speech-to-text with word-level timestamps, via WhisperX (faster-whisper backend). + +WhisperX gives us three things in one pipeline: + 1. Silero VAD to slice a 1h+ talk into speech regions (kills long-form hallucination). + 2. faster-whisper / CTranslate2 large-v3 transcription (Hindi, batched, on the GPU). + 3. wav2vec2 **forced alignment** -> per-word [start, end] accurate to ~tens of ms, + which is what makes "jump to the exact moment" reliable. + +Blackwell / RTX 50-series note: use compute_type="float16". int8 crashes with +CUBLAS_STATUS_NOT_SUPPORTED on sm_120; 12 GB VRAM fits float16 large-v3 comfortably. +""" +from __future__ import annotations + +# IMPORTANT: register CUDA/cuDNN DLL dirs BEFORE importing ctranslate2-backed libs. +from app import _cuda_bootstrap # noqa: F401 (side-effect import, must be first) + +import json +import wave +from pathlib import Path +from typing import Any, Dict, List, Optional + +import numpy as np + +from app.config import Config, get_config + +SAMPLE_RATE = 16000 + + +def read_wav_mono16k(path: str | Path) -> "np.ndarray": + """Read a 16 kHz mono PCM WAV into a float32 array in [-1, 1]. + + We always pre-normalize inputs to this exact format with ffmpeg (audio_utils), so we + can decode with the stdlib instead of depending on ffmpeg-on-PATH / torchaudio codecs. + """ + with wave.open(str(path), "rb") as w: + n_channels = w.getnchannels() + sampwidth = w.getsampwidth() + frames = w.readframes(w.getnframes()) + if sampwidth != 2: + raise ValueError(f"expected 16-bit PCM wav, got sample width {sampwidth} bytes") + audio = np.frombuffer(frames, dtype=np.int16) + if n_channels > 1: + audio = audio.reshape(-1, n_channels).mean(axis=1) + return np.ascontiguousarray(audio.astype(np.float32) / 32768.0) + + +class Transcriber: + """Lazy-loading wrapper around WhisperX. Load once, reuse for the whole run.""" + + def __init__(self, cfg: Optional[Config] = None): + self.cfg = cfg or get_config() + self._model = None + self._align_model = None + self._align_meta = None + self._align_lang: Optional[str] = None + self.device = self.cfg.asr["device"] + self.compute_type = self.cfg.asr["compute_type"] + + # ---- model loading ------------------------------------------------- + def _load_model(self): + if self._model is not None: + return self._model + import whisperx + + a = self.cfg.asr + asr_options = { + "beam_size": a["beam_size"], + "condition_on_previous_text": a["condition_on_previous_text"], + } + # whisperx.load_model signatures differ slightly across versions; pass what we can. + kwargs: Dict[str, Any] = dict( + device=self.device, + compute_type=self.compute_type, + language=a["language"], + asr_options=asr_options, + ) + if a.get("vad", True): + # Newer whisperx exposes vad_method; older uses default VAD. Try, then fall back. + try: + self._model = whisperx.load_model(a["model"], vad_method="silero", **kwargs) + except TypeError: + self._model = whisperx.load_model(a["model"], **kwargs) + else: + self._model = whisperx.load_model(a["model"], **kwargs) + return self._model + + def _load_align(self, language: str): + if self._align_model is not None and self._align_lang == language: + return self._align_model, self._align_meta + import whisperx + + model_name = self.cfg.asr.get("align_model") # None -> whisperx default for the lang + self._align_model, self._align_meta = whisperx.load_align_model( + language_code=language, device=self.device, model_name=model_name + ) + self._align_lang = language + return self._align_model, self._align_meta + + # ---- transcription ------------------------------------------------- + def transcribe(self, wav_path: str | Path) -> Dict[str, Any]: + """Full transcription + word alignment. Returns a transcript dict.""" + import whisperx + + wav_path = str(wav_path) + model = self._load_model() + audio = read_wav_mono16k(wav_path) + result = model.transcribe( + audio, batch_size=self.cfg.asr["batch_size"], language=self.cfg.asr["language"] + ) + language = result.get("language", self.cfg.asr["language"]) + + if self.cfg.asr.get("align", True) and result.get("segments"): + try: + align_model, meta = self._load_align(language) + result = whisperx.align( + result["segments"], + align_model, + meta, + audio, + self.device, + return_char_alignments=False, + ) + except Exception as exc: # noqa: BLE001 + print(f"[asr] alignment failed ({exc}); falling back to segment timestamps") + result["segments"] = _segments_to_pseudo_words(result["segments"]) + else: + result["segments"] = _segments_to_pseudo_words(result.get("segments", [])) + + duration = float(len(audio)) / 16000.0 + return { + "language": language, + "duration": duration, + "segments": _clean_segments(result.get("segments", [])), + } + + def transcribe_query(self, wav_path: str | Path) -> str: + """Fast path for a short query clip: just the Hindi text, no alignment.""" + import whisperx + + model = self._load_model() + audio = read_wav_mono16k(wav_path) + result = model.transcribe(audio, batch_size=self.cfg.asr["batch_size"], language=self.cfg.asr["language"]) + return " ".join(seg.get("text", "").strip() for seg in result.get("segments", [])).strip() + + +# ---- helpers ----------------------------------------------------------- +def _segments_to_pseudo_words(segments: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """When alignment is unavailable, synthesize one 'word' per segment spanning it. + + Search still works; jump-to-moment lands within a few seconds instead of ms. + """ + out = [] + for seg in segments: + if "start" not in seg or "end" not in seg: + continue + text = (seg.get("text") or "").strip() + seg = dict(seg) + seg["words"] = [{"word": text, "start": seg["start"], "end": seg["end"]}] + out.append(seg) + return out + + +def _clean_segments(segments: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Keep only well-formed words (numeric start/end). Backfill missing word times.""" + cleaned: List[Dict[str, Any]] = [] + for seg in segments: + words_in = seg.get("words") or [] + words_out: List[Dict[str, Any]] = [] + # Pass 1: keep words that already have numeric timing. + last_end = seg.get("start") + for w in words_in: + start = w.get("start") + end = w.get("end") + token = (w.get("word") or w.get("text") or "").strip() + if not token: + continue + if isinstance(start, (int, float)) and isinstance(end, (int, float)): + words_out.append({"word": token, "start": float(start), "end": float(end)}) + last_end = end + if not words_out: + # whole segment had no aligned words -> use the segment span + s = seg.get("start") + e = seg.get("end") + text = (seg.get("text") or "").strip() + if isinstance(s, (int, float)) and isinstance(e, (int, float)) and text: + words_out = [{"word": text, "start": float(s), "end": float(e)}] + if not words_out: + continue + cleaned.append( + { + "start": words_out[0]["start"], + "end": words_out[-1]["end"], + "text": (seg.get("text") or " ".join(w["word"] for w in words_out)).strip(), + "words": words_out, + } + ) + return cleaned + + +def save_transcript(transcript: Dict[str, Any], out_path: str | Path) -> None: + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with open(out_path, "w", encoding="utf-8") as fh: + json.dump(transcript, fh, ensure_ascii=False, indent=1) + + +def load_transcript(path: str | Path) -> Dict[str, Any]: + with open(path, "r", encoding="utf-8") as fh: + return json.load(fh) diff --git a/app/audio_source.py b/app/audio_source.py new file mode 100644 index 0000000000000000000000000000000000000000..207e91cfb2d75e9b41a4ab346f8055a0142d4c47 --- /dev/null +++ b/app/audio_source.py @@ -0,0 +1,36 @@ +"""Build the playable URL for a recording, for local serving or archive.org streaming. + +- ``local`` mode -> ``/audio/`` (served by StaticFiles from disk). +- ``archive`` mode -> ``https://archive.org/download//`` where the local + ``source_file`` (whose ``%5C`` are path separators) is remapped: drop the configured + local prefix, prepend the configured archive path prefix. Verified to range-seek (206). +""" +from __future__ import annotations + +from typing import Optional +from urllib.parse import quote + +from app.config import Config, get_config + + +def _archive_url(source_file: str, arch: dict) -> str: + rel = source_file.replace("%5C", "/") + prefix = (arch.get("strip_local_prefix") or "").strip("/") + if prefix and rel.startswith(prefix + "/"): + rel = rel[len(prefix) + 1:] + path_prefix = (arch.get("archive_path_prefix") or "").strip("/") + path = f"{path_prefix}/{rel}" if path_prefix else rel + base = (arch.get("base_url") or "https://archive.org/download").rstrip("/") + item = arch.get("item_id", "") + parts = ([item] if item else []) + path.split("/") + return base + "/" + "/".join(quote(seg) for seg in parts) + + +def audio_url(source_file: str, cfg: Optional[Config] = None) -> str: + cfg = cfg or get_config() + audio = cfg.get("audio", {}) or {} + if audio.get("mode", "local") == "archive": + return _archive_url(source_file, audio.get("archive", {}) or {}) + route = cfg.server["audio_route"] + encoded = "/".join(quote(seg) for seg in source_file.split("/")) + return f"{route}/{encoded}" diff --git a/app/audio_utils.py b/app/audio_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1ffd94ac44a9fcbf7249f82c3470c767ed8a7f42 --- /dev/null +++ b/app/audio_utils.py @@ -0,0 +1,56 @@ +"""ffmpeg helpers: locate the binary and normalize audio to 16 kHz mono WAV. + +Whisper/WhisperX expect 16 kHz mono; normalizing once up front also makes decoding of +exotic inputs (m4a/opus/wma/video containers) reliable. We always call ffmpeg by an +absolute path (missing-ffmpeg-on-PATH is the #1 cause of "the audio query does nothing"). +""" +from __future__ import annotations + +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Optional + +from app.config import get_config + +AUDIO_EXTS = { + ".mp3", ".wav", ".m4a", ".aac", ".flac", ".ogg", ".opus", + ".wma", ".mp4", ".webm", ".mkv", ".mov", ".3gp", +} + +_FFMPEG: Optional[str] = None + + +def find_ffmpeg() -> str: + """Locate ffmpeg: PATH first, then a bundled project ``bin/ffmpeg.exe``.""" + global _FFMPEG + if _FFMPEG: + return _FFMPEG + exe = shutil.which("ffmpeg") + if not exe: + binname = "ffmpeg.exe" if sys.platform == "win32" else "ffmpeg" + bundled = get_config().root / "bin" / binname + if bundled.exists(): + exe = str(bundled) + if not exe: + raise FileNotFoundError( + "ffmpeg not found. Install it (winget install Gyan.FFmpeg) or place ffmpeg.exe in ./bin/." + ) + _FFMPEG = exe + return exe + + +def normalize_to_wav(src: str | Path, dst: str | Path) -> Path: + """Decode any input to 16 kHz mono 16-bit PCM WAV at ``dst``.""" + dst = Path(dst) + dst.parent.mkdir(parents=True, exist_ok=True) + cmd = [ + find_ffmpeg(), "-y", "-hide_banner", "-loglevel", "error", + "-i", str(src), "-vn", "-ac", "1", "-ar", "16000", + "-c:a", "pcm_s16le", "-f", "wav", str(dst), + ] + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0 or not dst.exists(): + raise RuntimeError(f"ffmpeg failed for {src}: {proc.stderr.strip()[:500]}") + return dst diff --git a/app/chunking.py b/app/chunking.py new file mode 100644 index 0000000000000000000000000000000000000000..9202f7b624842e8a8532f4f41cff3e547da2429c --- /dev/null +++ b/app/chunking.py @@ -0,0 +1,124 @@ +"""Turn word-timestamped ASR output into searchable passages. + +Strategy (robust + precise): + * Each ASR segment is one "unit" (Whisper/VAD segments already break near natural + pauses, i.e. roughly sentence boundaries). We never split *across* a unit when + packing, so passages start/end on sentence-ish boundaries. + * Units are packed greedily into ~target_seconds windows (hard cap at max_seconds), + carrying ``overlap_sentences`` trailing units into the next window so a topic that + straddles a boundary is still findable. + * A passage's [start_ms, end_ms] always comes from its first/last *word* timestamps, + so playback seeks to the exact spoken moment. + * If a single segment is longer than max_seconds (rare; a long VAD region), it is + split by word timing into target-sized sub-passages. + +``split_sentences`` (Devanagari/danda-aware, via Indic NLP when available) is exposed for +reuse (e.g. snippet tidying / translation), with a regex fallback so it never hard-fails. +""" +from __future__ import annotations + +import re +from typing import Any, Dict, List, Optional + +from app.config import Config, get_config +from app.models import Passage + +# Sentence-final punctuation for Hindi (danda, double danda) + Latin. +_SENT_END = "।॥?!." +_SENT_SPLIT_RE = re.compile(r"(?<=[।॥?!])\s+|(?<=[.?!])\s+(?=[A-Zऀ-ॿ])") + + +def split_sentences(text: str) -> List[str]: + """Split Hindi/English text into sentences. Uses Indic NLP if installed, else regex.""" + text = (text or "").strip() + if not text: + return [] + try: + from indicnlp.tokenize import sentence_tokenize + + sents = sentence_tokenize.sentence_split(text, lang="hi") + sents = [s.strip() for s in sents if s.strip()] + if sents: + return sents + except Exception: # noqa: BLE001 — indicnlp optional / may not init; fall through + pass + parts = _SENT_SPLIT_RE.split(text) + return [p.strip() for p in parts if p.strip()] + + +def _split_long_unit(unit: Dict[str, Any], target_s: float) -> List[Dict[str, Any]]: + """Split a too-long unit into ~target_s sub-units by word timing.""" + words = unit["words"] + out: List[Dict[str, Any]] = [] + buf: List[Dict[str, Any]] = [] + start = words[0]["start"] + for w in words: + buf.append(w) + if w["end"] - start >= target_s: + out.append({"text": " ".join(x["word"] for x in buf), "words": buf}) + buf = [] + # next sub-unit starts at the next word + start = w["end"] + if buf: + out.append({"text": " ".join(x["word"] for x in buf), "words": buf}) + return out + + +def _segments_to_units(segments: List[Dict[str, Any]], max_s: float, target_s: float) -> List[Dict[str, Any]]: + units: List[Dict[str, Any]] = [] + for seg in segments: + words = seg.get("words") or [] + if not words: + continue + dur = words[-1]["end"] - words[0]["start"] + unit = {"text": (seg.get("text") or " ".join(w["word"] for w in words)).strip(), "words": words} + if dur > max_s: + units.extend(_split_long_unit(unit, target_s)) + else: + units.append(unit) + return units + + +def chunk_transcript( + transcript: Dict[str, Any], + recording_id: str, + source_file: str, + cfg: Optional[Config] = None, +) -> List[Passage]: + cfg = cfg or get_config() + c = cfg.chunking + target_s = float(c["target_seconds"]) + max_s = float(c["max_seconds"]) + overlap = int(c["overlap_sentences"]) + + units = _segments_to_units(transcript.get("segments", []), max_s, target_s) + if not units: + return [] + + passages: List[Passage] = [] + n = len(units) + i = 0 + while i < n: + buf_words: List[Dict[str, Any]] = [] + buf_texts: List[str] = [] + start_time = units[i]["words"][0]["start"] + j = i + while j < n: + u = units[j] + buf_words.extend(u["words"]) + buf_texts.append(u["text"]) + cur_end = u["words"][-1]["end"] + j += 1 + if cur_end - start_time >= target_s: + break # reached target; break on this unit boundary + text = " ".join(t for t in buf_texts if t).strip() + if buf_words and text: + passages.append(Passage.from_words(recording_id, source_file, buf_words, text)) + if j >= n: + break + i = max(j - overlap, i + 1) # carry `overlap` trailing units; always progress + return passages + + +def passages_text(passages: List[Passage]) -> List[str]: + return [p.hindi_text for p in passages] diff --git a/app/config.py b/app/config.py new file mode 100644 index 0000000000000000000000000000000000000000..5f3919ad77cf36b9fb7e30f178de35f43123ed25 --- /dev/null +++ b/app/config.py @@ -0,0 +1,190 @@ +"""Load and resolve configuration from ``config.yaml`` at the project root. + +Exposes a single ``get_config()`` returning a ``Config`` object with: + - dotted/dict access to every setting (``cfg.asr["model"]``, ``cfg["asr"]["model"]``) + - absolute, resolved paths under ``cfg.paths`` (audio_dir, work_dir, index_dir, models_dir) + - convenience properties for the LanceDB dir and the SQLite manifest path. + +Defaults are baked in, so the app still runs if a key (or the whole file) is missing. +""" +from __future__ import annotations + +import os +import sys +from functools import lru_cache +from pathlib import Path +from typing import Any, Dict + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +CONFIG_PATH = PROJECT_ROOT / "config.yaml" + +# Minimal defaults mirrored from config.yaml so the system degrades gracefully. +_DEFAULTS: Dict[str, Any] = { + "paths": { + "audio_dir": "data/audio", + "work_dir": "data/work", + "index_dir": "data/index", + "models_dir": "models", + }, + "asr": { + "model": "large-v3", + "device": "cuda", + "compute_type": "float16", + "language": "hi", + "beam_size": 5, + "batch_size": 8, + "vad": True, + "condition_on_previous_text": False, + "align": True, + "align_model": None, + }, + "chunking": { + "target_seconds": 45, + "max_seconds": 60, + "min_seconds": 8, + "overlap_sentences": 1, + }, + "embedding": { + "model": "BAAI/bge-m3", + "device": "cuda", + "use_fp16": True, + "dim": 1024, + "batch_size": 32, + "max_length": 1024, + }, + "reranker": { + "enabled": True, + "model": "BAAI/bge-reranker-v2-m3", + "device": "cuda", + "use_fp16": True, + "candidates": 80, + "batch_size": 16, + }, + "translation": { + "enabled": False, + "model": "ai4bharat/indictrans2-indic-en-1B", + "device": "cuda", + }, + "search": {"top_k": 20, "hybrid": True, "vector_weight": 0.7}, + "server": { + "host": "127.0.0.1", + "port": 8000, + "audio_route": "/audio", + }, + "audio": { + "mode": "local", + "archive": { + "base_url": "https://archive.org/download", + "item_id": "", + "strip_local_prefix": "", + "archive_path_prefix": "", + }, + }, +} + + +def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]: + out = dict(base) + for key, val in (override or {}).items(): + if isinstance(val, dict) and isinstance(out.get(key), dict): + out[key] = _deep_merge(out[key], val) + else: + out[key] = val + return out + + +class Config: + """Thin wrapper over the merged config dict with resolved absolute paths.""" + + def __init__(self, data: Dict[str, Any]): + self._data = data + self.root = PROJECT_ROOT + # Resolve every path to an absolute Path and ensure the directory exists. + self.paths = { + name: self._resolve(rel) for name, rel in data["paths"].items() + } + + def _resolve(self, rel: str) -> Path: + p = Path(rel) + if not p.is_absolute(): + p = PROJECT_ROOT / p + return p.resolve() + + def ensure_dirs(self) -> None: + for p in self.paths.values(): + p.mkdir(parents=True, exist_ok=True) + + # ---- access helpers ------------------------------------------------- + def __getitem__(self, key: str) -> Any: + return self._data[key] + + def __getattr__(self, name: str) -> Any: + # Only called for attributes not found normally (so self._data etc. are safe). + try: + return self._data[name] + except KeyError as exc: # pragma: no cover + raise AttributeError(name) from exc + + def get(self, key: str, default: Any = None) -> Any: + return self._data.get(key, default) + + # ---- convenience paths --------------------------------------------- + @property + def lancedb_dir(self) -> Path: + return self.paths["index_dir"] / "lancedb" + + @property + def manifest_path(self) -> Path: + return self.paths["index_dir"] / "manifest.sqlite" + + @property + def audio_dir(self) -> Path: + return self.paths["audio_dir"] + + @property + def work_dir(self) -> Path: + return self.paths["work_dir"] + + @property + def models_dir(self) -> Path: + return self.paths["models_dir"] + + def as_dict(self) -> Dict[str, Any]: + return dict(self._data) + + +@lru_cache(maxsize=1) +def get_config() -> Config: + data = dict(_DEFAULTS) + if CONFIG_PATH.exists(): + try: + import yaml # local import: only needed when a config file is present + + with open(CONFIG_PATH, "r", encoding="utf-8") as fh: + loaded = yaml.safe_load(fh) or {} + data = _deep_merge(_DEFAULTS, loaded) + except Exception as exc: # noqa: BLE001 + print(f"[config] WARNING: could not read {CONFIG_PATH}: {exc}; using defaults") + # Deployment can flip audio mode without editing the file (copy audio to avoid + # mutating the shared _DEFAULTS object). + env_mode = os.environ.get("SEARCHAUDIO_AUDIO_MODE") + if env_mode: + audio = dict(data.get("audio", {}) or {}) + audio["mode"] = env_mode + data["audio"] = audio + cfg = Config(data) + # Point the HF cache at our local models dir so everything stays in the project + # and works fully offline (after a one-time download). + os.environ.setdefault("HF_HOME", str(cfg.models_dir)) + os.environ.setdefault("HF_HUB_CACHE", str(cfg.models_dir / "hub")) + # Windows without Developer Mode/admin can't create symlinks (WinError 1314 during + # model downloads). Tell huggingface_hub to copy files instead. + if sys.platform == "win32": + os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS", "1") + os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1") + # Put the bundled ffmpeg on PATH so any library that shells out to `ffmpeg` by name + # (e.g. whisperx.load_audio) can find it, even though it lives in ./bin. + bin_dir = PROJECT_ROOT / "bin" + if bin_dir.is_dir() and str(bin_dir) not in os.environ.get("PATH", ""): + os.environ["PATH"] = str(bin_dir) + os.pathsep + os.environ.get("PATH", "") + return cfg diff --git a/app/embed.py b/app/embed.py new file mode 100644 index 0000000000000000000000000000000000000000..ad4162e44973a0c40d9b1102fa708b383ef1e5db --- /dev/null +++ b/app/embed.py @@ -0,0 +1,53 @@ +"""Dense text embeddings via BAAI/bge-m3. + +bge-m3 maps Hindi passages and an English query into the **same** vector space, so an +English sentence retrieves Hindi content natively — no query-time translation. It needs +no "query:"/"passage:" prefix (symmetric), and its dense vectors are unit-normalized, so +cosine == dot product. +""" +from __future__ import annotations + +from typing import List, Optional + +import numpy as np + +from app.config import Config, get_config + + +class Embedder: + def __init__(self, cfg: Optional[Config] = None): + self.cfg = cfg or get_config() + self._model = None + + def _load(self): + if self._model is not None: + return self._model + from FlagEmbedding import BGEM3FlagModel + + e = self.cfg.embedding + # FlagEmbedding auto-selects CUDA when available; use_fp16 ~halves VRAM and ~2x speed. + self._model = BGEM3FlagModel(e["model"], use_fp16=bool(e["use_fp16"])) + return self._model + + def embed_passages(self, texts: List[str]) -> np.ndarray: + """Return an (N, dim) float32 array of dense embeddings.""" + if not texts: + return np.zeros((0, self.cfg.embedding["dim"]), dtype=np.float32) + model = self._load() + e = self.cfg.embedding + out = model.encode( + texts, + batch_size=e["batch_size"], + max_length=e["max_length"], + return_dense=True, + return_sparse=False, + return_colbert_vecs=False, + ) + vecs = np.asarray(out["dense_vecs"], dtype=np.float32) + if vecs.ndim == 1: + vecs = vecs.reshape(1, -1) + return np.ascontiguousarray(vecs) + + def embed_query(self, text: str) -> np.ndarray: + """Return a single (dim,) float32 vector for a query string.""" + return self.embed_passages([text])[0] diff --git a/app/ingest.py b/app/ingest.py new file mode 100644 index 0000000000000000000000000000000000000000..c1d7ae62597616a2f4e7bacb55ab3a95963bd938 --- /dev/null +++ b/app/ingest.py @@ -0,0 +1,164 @@ +"""End-to-end ingestion: audio folder -> searchable index. Resumable & incremental. + +Per file: normalize (ffmpeg) -> transcribe+align (WhisperX) -> chunk -> [translate] + -> embed (bge-m3) -> upsert (LanceDB). + +Resumability: the SQLite manifest records each file's furthest-completed status; the +normalized WAV and transcript JSON are cached on disk. Re-running skips finished work and +continues from where it stopped. One bad file is marked ``failed`` and never blocks the +rest. Incremental ingest is just re-running ``run_ingest`` — only new content hashes are +processed. +""" +from __future__ import annotations + +import traceback +from pathlib import Path +from typing import Optional + +from app.asr import Transcriber, load_transcript, save_transcript +from app.audio_utils import AUDIO_EXTS, normalize_to_wav +from app.chunking import chunk_transcript +from app.config import Config, get_config +from app.manifest import Manifest +from app.models import Passage, file_content_hash, make_recording_id +from app.store import Store + + +class Pipeline: + """Holds the heavy models so they load once and are reused across all files.""" + + def __init__(self, cfg: Optional[Config] = None): + self.cfg = cfg or get_config() + self.cfg.ensure_dirs() + self.store = Store(self.cfg) + self._transcriber: Optional[Transcriber] = None + self._embedder = None + self._translator = None + + @property + def transcriber(self) -> Transcriber: + if self._transcriber is None: + self._transcriber = Transcriber(self.cfg) + return self._transcriber + + @property + def embedder(self): + if self._embedder is None: + from app.embed import Embedder + + self._embedder = Embedder(self.cfg) + return self._embedder + + @property + def translator(self): + if self._translator is None: + from app.translate import Translator + + self._translator = Translator(self.cfg) + return self._translator + + # ---- per-file processing ------------------------------------------- + def process(self, content_hash: str, abs_path: str, rel_path: str, manifest: Manifest) -> int: + rid = make_recording_id(content_hash) + work = self.cfg.work_dir + wav_path = work / f"{rid}.wav" + transcript_path = work / f"{rid}.json" + + # 1. normalize (cache) + if not wav_path.exists(): + normalize_to_wav(abs_path, wav_path) + manifest.set_status(content_hash, "normalized") + + # 2. transcribe + align (cache the transcript JSON) + if transcript_path.exists(): + transcript = load_transcript(transcript_path) + else: + transcript = self.transcriber.transcribe(wav_path) + save_transcript(transcript, transcript_path) + manifest.set_status(content_hash, "transcribed", duration_s=transcript.get("duration")) + + # 3. chunk + passages = chunk_transcript(transcript, rid, rel_path, self.cfg) + manifest.set_status(content_hash, "chunked", n_chunks=len(passages)) + if not passages: + manifest.set_status(content_hash, "indexed", n_chunks=0) + return 0 + + # 4. optional translation for English snippets + if self.cfg.translation["enabled"]: + try: + glosses = self.translator.translate([p.hindi_text for p in passages]) + for p, g in zip(passages, glosses): + p.english_gloss = g + except Exception as exc: # noqa: BLE001 — translation is non-essential + print(f"[ingest] translation skipped for {rel_path}: {exc}") + + # 5. embed + vectors = self.embedder.embed_passages([p.hindi_text for p in passages]) + manifest.set_status(content_hash, "embedded") + + # 6. index (idempotent: clear any prior rows for this recording first) + self.store.delete_recording(rid) + self.store.upsert_passages(passages, vectors) + manifest.set_status(content_hash, "indexed", n_chunks=len(passages)) + return len(passages) + + +def scan_audio_dir(manifest: Manifest, cfg: Optional[Config] = None) -> int: + """Register any new audio files in the manifest. Returns count of newly added.""" + cfg = cfg or get_config() + audio_dir = cfg.audio_dir + audio_dir.mkdir(parents=True, exist_ok=True) + new = 0 + for path in sorted(audio_dir.rglob("*")): + if not path.is_file() or path.suffix.lower() not in AUDIO_EXTS: + continue + rel = path.relative_to(audio_dir).as_posix() + h = file_content_hash(str(path)) + if manifest.upsert_file(h, rel, str(path)): + new += 1 + return new + + +def run_ingest( + cfg: Optional[Config] = None, + limit: Optional[int] = None, + optimize_every: int = 25, + include_failed: bool = False, +) -> dict: + """Scan + process all pending files. Safe to interrupt and re-run.""" + cfg = cfg or get_config() + cfg.ensure_dirs() + manifest = Manifest(cfg.manifest_path) + pipe = Pipeline(cfg) + + added = scan_audio_dir(manifest, cfg) + pending = manifest.pending(include_failed=include_failed) + if limit: + pending = pending[:limit] + print(f"[ingest] {added} new file(s) found; {len(pending)} to process") + + done = failed = 0 + for idx, row in enumerate(pending, 1): + h, rel, ap = row["content_hash"], row["rel_path"], row["abs_path"] + print(f"[ingest] ({idx}/{len(pending)}) {rel}") + try: + n = pipe.process(h, ap, rel, manifest) + done += 1 + print(f" -> indexed {n} passage(s)") + except Exception as exc: # noqa: BLE001 — isolate per-file failures + failed += 1 + manifest.set_status(h, "failed", error=f"{exc}") + print(f" -> FAILED: {exc}") + traceback.print_exc() + if idx % optimize_every == 0: + pipe.store.optimize() + + pipe.store.optimize() + if cfg.search["hybrid"]: + pipe.store.ensure_fts_index() + counts = manifest.counts() + manifest.close() + summary = {"new_files": added, "processed": done, "failed": failed, "manifest": counts} + print(f"[ingest] done: {summary}") + return summary diff --git a/app/manifest.py b/app/manifest.py new file mode 100644 index 0000000000000000000000000000000000000000..d66ff3eeb6820b4bd521ee05833eb4fab68b8c20 --- /dev/null +++ b/app/manifest.py @@ -0,0 +1,139 @@ +"""SQLite manifest — the source of truth for what has been ingested. + +Keyed on the file's **content hash** (not its path), so moving or renaming a recording +never triggers a re-transcribe. Each file walks through a status ladder; the bulk +ingest loop simply re-processes anything whose status is below ``indexed`` (and not +``failed``), which makes the whole multi-day run trivially resumable: kill it any time, +restart, and it continues where it left off. +""" +from __future__ import annotations + +import sqlite3 +import time +from pathlib import Path +from typing import Dict, List, Optional + +# Ordered ladder. A file is "done" only at INDEXED. FAILED is terminal (until retried). +STATUS_LADDER = ["pending", "normalized", "transcribed", "chunked", "embedded", "indexed"] +STATUS_FAILED = "failed" +STATUS_DONE = "indexed" + + +def _now() -> str: + return time.strftime("%Y-%m-%dT%H:%M:%S") + + +class Manifest: + def __init__(self, db_path: str | Path): + self.db_path = str(db_path) + Path(self.db_path).parent.mkdir(parents=True, exist_ok=True) + self.conn = sqlite3.connect(self.db_path) + self.conn.row_factory = sqlite3.Row + self.conn.execute("PRAGMA journal_mode=WAL;") + self._init_db() + + def _init_db(self) -> None: + self.conn.execute( + """ + CREATE TABLE IF NOT EXISTS files ( + content_hash TEXT PRIMARY KEY, + rel_path TEXT NOT NULL, + abs_path TEXT NOT NULL, + duration_s REAL, + status TEXT NOT NULL DEFAULT 'pending', + error TEXT, + n_chunks INTEGER DEFAULT 0, + added_at TEXT, + updated_at TEXT + ) + """ + ) + self.conn.execute("CREATE INDEX IF NOT EXISTS idx_status ON files(status);") + self.conn.commit() + + # ---- writes --------------------------------------------------------- + def upsert_file(self, content_hash: str, rel_path: str, abs_path: str) -> bool: + """Register a file. Returns True if it's new (caller may process it).""" + row = self.get(content_hash) + now = _now() + if row is None: + self.conn.execute( + "INSERT INTO files (content_hash, rel_path, abs_path, status, added_at, updated_at) " + "VALUES (?, ?, ?, 'pending', ?, ?)", + (content_hash, rel_path, abs_path, now, now), + ) + self.conn.commit() + return True + # Known file: keep status, but refresh path if it moved. + if row["abs_path"] != abs_path or row["rel_path"] != rel_path: + self.conn.execute( + "UPDATE files SET rel_path=?, abs_path=?, updated_at=? WHERE content_hash=?", + (rel_path, abs_path, now, content_hash), + ) + self.conn.commit() + return False + + def set_status( + self, + content_hash: str, + status: str, + *, + error: Optional[str] = None, + duration_s: Optional[float] = None, + n_chunks: Optional[int] = None, + ) -> None: + sets = ["status=?", "updated_at=?"] + vals: List[object] = [status, _now()] + if error is not None: + sets.append("error=?") + vals.append(error) + if duration_s is not None: + sets.append("duration_s=?") + vals.append(duration_s) + if n_chunks is not None: + sets.append("n_chunks=?") + vals.append(n_chunks) + vals.append(content_hash) + self.conn.execute(f"UPDATE files SET {', '.join(sets)} WHERE content_hash=?", vals) + self.conn.commit() + + def remove(self, content_hash: str) -> None: + self.conn.execute("DELETE FROM files WHERE content_hash=?", (content_hash,)) + self.conn.commit() + + # ---- reads ---------------------------------------------------------- + def get(self, content_hash: str) -> Optional[sqlite3.Row]: + cur = self.conn.execute("SELECT * FROM files WHERE content_hash=?", (content_hash,)) + return cur.fetchone() + + def pending(self, include_failed: bool = False) -> List[sqlite3.Row]: + """Files that still need work (status below 'indexed').""" + if include_failed: + cur = self.conn.execute( + "SELECT * FROM files WHERE status != ? ORDER BY added_at", (STATUS_DONE,) + ) + else: + cur = self.conn.execute( + "SELECT * FROM files WHERE status != ? AND status != ? ORDER BY added_at", + (STATUS_DONE, STATUS_FAILED), + ) + return cur.fetchall() + + def all(self) -> List[sqlite3.Row]: + return self.conn.execute("SELECT * FROM files ORDER BY added_at").fetchall() + + def known_hashes(self) -> set: + return {r["content_hash"] for r in self.conn.execute("SELECT content_hash FROM files")} + + def counts(self) -> Dict[str, int]: + cur = self.conn.execute("SELECT status, COUNT(*) c FROM files GROUP BY status") + return {r["status"]: r["c"] for r in cur.fetchall()} + + def close(self) -> None: + self.conn.close() + + def __enter__(self) -> "Manifest": + return self + + def __exit__(self, *exc) -> None: + self.close() diff --git a/app/models.py b/app/models.py new file mode 100644 index 0000000000000000000000000000000000000000..5690785f3a0f14bd04ac9cb49ae62d93f99676bd --- /dev/null +++ b/app/models.py @@ -0,0 +1,86 @@ +"""Shared data structures passed between pipeline stages. + +A ``Passage`` is one searchable unit: a sentence-packed window of transcript with a +precise [start_ms, end_ms] span inside one recording. It is what we embed, index, and +return to the UI (so the player can seek straight to ``start_ms``). +""" +from __future__ import annotations + +import hashlib +from dataclasses import asdict, dataclass, field +from typing import Any, Dict, List, Optional + + +def make_recording_id(content_hash: str) -> str: + """Short, stable id for a recording, derived from its content hash.""" + return content_hash[:16] + + +def make_passage_id(recording_id: str, start_ms: int) -> str: + """Stable id for a passage (unique within the corpus, deterministic on re-ingest).""" + return f"{recording_id}:{start_ms:09d}" + + +@dataclass +class Passage: + id: str + recording_id: str + source_file: str # path relative to the audio_dir (used to build the audio URL) + start_ms: int + end_ms: int + hindi_text: str + n_words: int = 0 + english_gloss: str = "" # filled only when translation is enabled + + def to_row(self) -> Dict[str, Any]: + """Row dict for LanceDB (vector is added separately at index time).""" + return asdict(self) + + @staticmethod + def from_words( + recording_id: str, + source_file: str, + words: List[Dict[str, Any]], + text: str, + ) -> "Passage": + start_ms = int(round(words[0]["start"] * 1000)) + end_ms = int(round(words[-1]["end"] * 1000)) + return Passage( + id=make_passage_id(recording_id, start_ms), + recording_id=recording_id, + source_file=source_file, + start_ms=start_ms, + end_ms=end_ms, + hindi_text=text.strip(), + n_words=len(words), + ) + + +@dataclass +class SearchResult: + recording_id: str + source_file: str + start_ms: int + end_ms: int + hindi_text: str + english_gloss: str + score: float + rerank_score: Optional[float] = None + + @property + def start_seconds(self) -> float: + return self.start_ms / 1000.0 + + def to_dict(self) -> Dict[str, Any]: + d = asdict(self) + d["start_seconds"] = round(self.start_seconds, 2) + return d + + +def file_content_hash(path: str, chunk_size: int = 1 << 20) -> str: + """SHA-256 of the file's bytes. Keys the manifest so renames don't re-transcribe.""" + h = hashlib.sha256() + with open(path, "rb") as fh: + for block in iter(lambda: fh.read(chunk_size), b""): + h.update(block) + return h.hexdigest() diff --git a/app/query.py b/app/query.py new file mode 100644 index 0000000000000000000000000000000000000000..fcbe38e44df533f8b409d017aacb59ac17394914 --- /dev/null +++ b/app/query.py @@ -0,0 +1,112 @@ +"""Query side: text or Hindi-audio -> ranked moments. + + text query -> bge-m3 embed -> LanceDB (dense-dominant hybrid) -> bge-reranker -> top-k + audio query -> ffmpeg -> faster-whisper(hi) -> Hindi text -> (same as above) + +An English query matches Hindi passages natively in bge-m3's shared space, so there is no +query-time translation. The reranker reads (query, Hindi passage) pairs together and is +the main precision lever. +""" +from __future__ import annotations + +import tempfile +from pathlib import Path +from typing import List, Optional, Tuple + +from app.config import Config, get_config +from app.models import SearchResult +from app.store import Store + + +class Searcher: + def __init__(self, cfg: Optional[Config] = None): + self.cfg = cfg or get_config() + self.store = Store(self.cfg) + self._embedder = None + self._reranker = None + self._transcriber = None + + @property + def embedder(self): + if self._embedder is None: + from app.embed import Embedder + + self._embedder = Embedder(self.cfg) + return self._embedder + + @property + def reranker(self): + if self._reranker is None: + from app.rerank import Reranker + + self._reranker = Reranker(self.cfg) + return self._reranker + + @property + def transcriber(self): + if self._transcriber is None: + from app.asr import Transcriber + + self._transcriber = Transcriber(self.cfg) + return self._transcriber + + # ---- text search ---------------------------------------------------- + def search_text(self, query: str, top_k: Optional[int] = None) -> List[SearchResult]: + query = (query or "").strip() + if not query: + return [] + top_k = top_k or self.cfg.search["top_k"] + use_rerank = self.reranker.enabled + pool = max(self.cfg.reranker["candidates"], top_k) if use_rerank else top_k + + qvec = self.embedder.embed_query(query) + rows = self.store.search( + qvec, + query_text=query, + k=pool, + hybrid=bool(self.cfg.search["hybrid"]), + vector_weight=float(self.cfg.search["vector_weight"]), + candidate_pool=max(pool, 80), + ) + if not rows: + return [] + + if use_rerank: + scores = self.reranker.scores(query, [r["hindi_text"] for r in rows]) + for r, s in zip(rows, scores): + r["rerank_score"] = s + rows.sort(key=lambda r: r.get("rerank_score", 0.0), reverse=True) + + rows = rows[:top_k] + return [self._to_result(r) for r in rows] + + # ---- audio search --------------------------------------------------- + def search_audio(self, audio_path: str | Path, top_k: Optional[int] = None) -> Tuple[str, List[SearchResult]]: + """Transcribe a Hindi clip, then search. Returns (recognized_text, results).""" + from app.audio_utils import normalize_to_wav + + with tempfile.TemporaryDirectory() as tmp: + wav = Path(tmp) / "query.wav" + normalize_to_wav(audio_path, wav) + text = self.transcriber.transcribe_query(wav) + results = self.search_text(text, top_k=top_k) if text.strip() else [] + return text, results + + # ---- helpers -------------------------------------------------------- + @staticmethod + def _to_result(r: dict) -> SearchResult: + # Prefer the reranker score, else the dense cosine (nice 0..1 for display); + # fused RRF score is only a last resort (tiny values, ranking-only). + score = r.get("rerank_score") + if score is None: + score = r.get("score", r.get("fused_score", 0.0)) + return SearchResult( + recording_id=r["recording_id"], + source_file=r["source_file"], + start_ms=int(r["start_ms"]), + end_ms=int(r["end_ms"]), + hindi_text=r.get("hindi_text", ""), + english_gloss=r.get("english_gloss", ""), + score=float(score), + rerank_score=(float(r["rerank_score"]) if "rerank_score" in r else None), + ) diff --git a/app/rerank.py b/app/rerank.py new file mode 100644 index 0000000000000000000000000000000000000000..07f848d15f9099b64f5f9d6fae2e203bafe50b4c --- /dev/null +++ b/app/rerank.py @@ -0,0 +1,44 @@ +"""Cross-encoder reranking via BAAI/bge-reranker-v2-m3. + +Dense retrieval is great for recall but ranks by a single vector dot-product. A +cross-encoder reads (query, passage) **together** and scores true relevance — the single +biggest precision lever for cross-lingual search. We only score the top ~80 candidates, +so it's cheap. ``normalize=True`` squashes scores to 0..1 (sigmoid) for nice display. +""" +from __future__ import annotations + +from typing import List, Optional + +from app.config import Config, get_config + + +class Reranker: + def __init__(self, cfg: Optional[Config] = None): + self.cfg = cfg or get_config() + self._model = None + + @property + def enabled(self) -> bool: + return bool(self.cfg.reranker["enabled"]) + + def _load(self): + if self._model is not None: + return self._model + from FlagEmbedding import FlagReranker + + r = self.cfg.reranker + self._model = FlagReranker(r["model"], use_fp16=bool(r["use_fp16"])) + return self._model + + def scores(self, query: str, passages: List[str]) -> List[float]: + """Relevance score (0..1) for each passage against the query.""" + if not passages: + return [] + model = self._load() + pairs = [[query, p] for p in passages] + out = model.compute_score( + pairs, batch_size=self.cfg.reranker["batch_size"], normalize=True + ) + if isinstance(out, (int, float)): + return [float(out)] + return [float(x) for x in out] diff --git a/app/server.py b/app/server.py new file mode 100644 index 0000000000000000000000000000000000000000..effd6ead9767251b7ad4b96fe080e64f38f0a019 --- /dev/null +++ b/app/server.py @@ -0,0 +1,135 @@ +"""Local FastAPI app: search box + clickable, seekable results. + +Endpoints + GET / -> the search page + POST /api/search -> {query, top_k} -> ranked moments (JSON) + POST /api/query_audio -> multipart audio file -> {text, results} + GET /api/stats -> index / manifest stats + GET {audio_route}/... -> the original recordings, served with HTTP Range so the +