diff --git a/desktop/capability-catalog.json b/desktop/capability-catalog.json index 8297ae0d..ee0da1e1 100644 --- a/desktop/capability-catalog.json +++ b/desktop/capability-catalog.json @@ -5,7 +5,7 @@ "extra": "speech", "modality": "speech", "label": "Speech search", - "description": "Transcribe and search spoken words with timestamps.", + "description": "Transcribe spoken words once, build searchable segments, and retrieve them with semantic and exact keyword matching.", "models": [ { "cache_key": "models--Qwen--Qwen3-Embedding-0.6B/snapshots/97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3/model.safetensors", diff --git a/src/vidxp/benchmarks/hirest.py b/src/vidxp/benchmarks/hirest.py index dbe97985..61841778 100644 --- a/src/vidxp/benchmarks/hirest.py +++ b/src/vidxp/benchmarks/hirest.py @@ -522,6 +522,12 @@ def run_hirest( "dialogue_words_per_phrase": ( speech_config(config).words_per_phrase ), + "dialogue_segmentation_mode": ( + speech_config(config).segmentation_mode + ), + "dialogue_window_stride_words": ( + speech_config(config).window_stride_words + ), "segment_word_timestamps": ( "linear_interpolation_within_srt_cue" ), diff --git a/src/vidxp/capabilities/search.py b/src/vidxp/capabilities/search.py index 9c57ddc3..316ac1dc 100644 --- a/src/vidxp/capabilities/search.py +++ b/src/vidxp/capabilities/search.py @@ -25,6 +25,10 @@ "representation", "window_index", "activation_index", + "match_kind", + "word_start", + "word_end", + "segmentation_mode", } ) @@ -60,7 +64,7 @@ def stable_query_id( return f"{modality}:{digest}" -def _to_hits( +def hits_from_rows( modality: str, rows: list[dict[str, Any]], required_metadata: frozenset[str], @@ -143,7 +147,7 @@ def search_embeddings( query_id=query_id or stable_query_id(query, modality, config), query=query, modality=modality, - hits=_to_hits(modality, rows, required_metadata), + hits=hits_from_rows(modality, rows, required_metadata), ) diff --git a/src/vidxp/capabilities/speech/__init__.py b/src/vidxp/capabilities/speech/__init__.py index e21af745..4b8a0c54 100644 --- a/src/vidxp/capabilities/speech/__init__.py +++ b/src/vidxp/capabilities/speech/__init__.py @@ -1 +1 @@ -"""Timestamped speech transcription and semantic search.""" +"""Timed speech transcripts with semantic and keyword search.""" diff --git a/src/vidxp/capabilities/speech/config.py b/src/vidxp/capabilities/speech/config.py index 618df26e..30137cd1 100644 --- a/src/vidxp/capabilities/speech/config.py +++ b/src/vidxp/capabilities/speech/config.py @@ -1,13 +1,30 @@ from __future__ import annotations +from typing import Literal + from pydantic import Field from vidxp.capabilities.contracts import CapabilityConfig from vidxp.core.contracts import IndexConfig +SegmentationMode = Literal[ + "fixed_words", + "overlapping_windows", + "sentence", +] + class SpeechConfig(CapabilityConfig): + """Speech indexing and search settings. + + Changing segmentation fields changes ``IndexConfig.fingerprint()`` and + therefore requires a new index generation rather than silently rewriting + existing segment IDs. + """ + words_per_phrase: int = Field(default=5, gt=0) + window_stride_words: int = Field(default=2, gt=0) + segmentation_mode: SegmentationMode = "fixed_words" embedding_batch_size: int = Field(default=128, gt=0) transcription_batch_size: int = Field(default=16, gt=0) normalize_embeddings: bool = True diff --git a/src/vidxp/capabilities/speech/definition.py b/src/vidxp/capabilities/speech/definition.py index 63771bca..48dc894f 100644 --- a/src/vidxp/capabilities/speech/definition.py +++ b/src/vidxp/capabilities/speech/definition.py @@ -84,7 +84,10 @@ def model_manifest( DEFINITION = CapabilityDefinition( name="speech", label="Speech search", - description="Transcribe and search spoken words with timestamps.", + description=( + "Transcribe spoken words once, build searchable segments, and " + "retrieve them with semantic and exact keyword matching." + ), extra="speech", config_model=SpeechConfig, collection_name="speech", diff --git a/src/vidxp/capabilities/speech/indexing.py b/src/vidxp/capabilities/speech/indexing.py index c8ee42d6..7d2cd024 100644 --- a/src/vidxp/capabilities/speech/indexing.py +++ b/src/vidxp/capabilities/speech/indexing.py @@ -1,15 +1,22 @@ from __future__ import annotations -from dataclasses import dataclass from pathlib import Path from typing import Any, Mapping, Sequence -from vidxp.capabilities.speech.config import speech_config +from vidxp.capabilities.speech.config import SegmentationMode, speech_config from vidxp.capabilities.speech.models import get_embedder, get_whisper_model +from vidxp.capabilities.speech.segmentation import ( + DialoguePhrase, + build_dialogue_phrases_from_words, +) from vidxp.capabilities.speech.specs import ( FASTER_WHISPER_MODEL, QWEN3_EMBEDDING_MODEL, ) +from vidxp.capabilities.speech.transcript import ( + flatten_transcript_words, + save_transcript, +) from vidxp.core.contracts import ( CancellationToken, IndexConfig, @@ -21,95 +28,21 @@ from vidxp.ports import IndexStore, ModelRuntimePort -@dataclass(frozen=True) -class DialoguePhrase: - phrase_id: int - text: str - start: float - end: float - - -def _valid_interval(start: Any, end: Any, label: str) -> tuple[float, float]: - start_value = float(start) - end_value = float(end) - if start_value < 0 or end_value <= start_value: - raise ValueError( - f"{label} must have a non-negative, non-zero interval; " - f"received [{start_value}, {end_value}]." - ) - return start_value, end_value - - def build_dialogue_phrases( segments: Sequence[Mapping[str, Any]], *, words_per_phrase: int, + segmentation_mode: SegmentationMode = "fixed_words", + window_stride_words: int = 2, ) -> list[DialoguePhrase]: - if words_per_phrase <= 0: - raise ValueError("words_per_phrase must be greater than zero.") - phrases: list[DialoguePhrase] = [] - for segment_index, segment in enumerate(segments): - words = segment.get("words") or [] - if words: - timestamped = [ - word - for word in words - if str(word.get("word", word.get("text", ""))).strip() - and word.get("start") is not None - and word.get("end") is not None - ] - for offset in range(0, len(timestamped), words_per_phrase): - group = timestamped[offset:offset + words_per_phrase] - if not group: - continue - start, end = _valid_interval( - group[0]["start"], - group[-1]["end"], - f"Transcript word group in segment {segment_index}", - ) - text = " ".join( - str(word.get("word", word.get("text", ""))).strip() - for word in group - ) - phrases.append( - DialoguePhrase( - phrase_id=len(phrases), - text=text, - start=start, - end=end, - ) - ) - continue + """Build searchable phrases from Whisper or supplied transcript segments.""" - text = str(segment.get("text", "")).strip() - if not text: - continue - if segment.get("start") is None or segment.get("end") is None: - raise ValueError( - f"Transcript segment {segment_index} lacks start/end timestamps." - ) - start, end = _valid_interval( - segment["start"], - segment["end"], - f"Transcript segment {segment_index}", - ) - tokens = text.split() - duration = end - start - for offset in range(0, len(tokens), words_per_phrase): - group = tokens[offset:offset + words_per_phrase] - group_start = start + duration * offset / len(tokens) - group_end = start + duration * ( - offset + len(group) - ) / len(tokens) - phrases.append( - DialoguePhrase( - phrase_id=len(phrases), - text=" ".join(group), - start=group_start, - end=group_end, - ) - ) - return phrases + return build_dialogue_phrases_from_words( + flatten_transcript_words(segments), + segmentation_mode=segmentation_mode, + words_per_phrase=words_per_phrase, + window_stride_words=window_stride_words, + ) def transcribe_video( @@ -176,7 +109,7 @@ def transcribe_video( def _speech_records( - phrases, + phrases: Sequence[DialoguePhrase], vectors, config: IndexConfig, ) -> list[StorageRecord]: @@ -186,7 +119,7 @@ def _speech_records( config.run_id, str(config.video_id), "speech", - f"p{phrase.phrase_id:08d}", + phrase.local_id, generation_id=config.generation_id, ) records.append( @@ -200,6 +133,9 @@ def _speech_records( "text": phrase.text, "start": phrase.start, "end": phrase.end, + "word_start": phrase.word_start, + "word_end": phrase.word_end, + "segmentation_mode": phrase.segmentation_mode, }, ) ) @@ -235,12 +171,27 @@ def index_speech( progress=progress, ) - phrases = build_dialogue_phrases( - segments, + words = flatten_transcript_words(segments) + if words: + save_transcript( + config.run_directory, + words, + language=language, + ) + + phrases = build_dialogue_phrases_from_words( + words, + segmentation_mode=settings.segmentation_mode, words_per_phrase=settings.words_per_phrase, + window_stride_words=settings.window_stride_words, ) if not phrases: - return {"dialogue_phrases": 0, "language": language} + return { + "dialogue_phrases": 0, + "transcript_words": len(words), + "segmentation_mode": settings.segmentation_mode, + "language": language, + } report_progress( progress, @@ -280,4 +231,9 @@ def index_speech( stored, len(phrases), ) - return {"dialogue_phrases": stored, "language": language} + return { + "dialogue_phrases": stored, + "transcript_words": len(words), + "segmentation_mode": settings.segmentation_mode, + "language": language, + } diff --git a/src/vidxp/capabilities/speech/operations.py b/src/vidxp/capabilities/speech/operations.py index b5b5431c..af9380e0 100644 --- a/src/vidxp/capabilities/speech/operations.py +++ b/src/vidxp/capabilities/speech/operations.py @@ -1,5 +1,6 @@ from __future__ import annotations +import re from typing import Any, Mapping from vidxp.capabilities.contracts import ( @@ -11,7 +12,7 @@ from vidxp.capabilities.speech.indexing import index_speech from vidxp.capabilities.speech.models import get_embedder from vidxp.capabilities.schemas import SearchInput, SearchResult -from vidxp.capabilities.search import search_embeddings +from vidxp.capabilities.search import hits_from_rows, stable_query_id from vidxp.core.contracts import ( CancellationToken, IndexConfig, @@ -36,6 +37,12 @@ } ) +_TOKEN_PATTERN = re.compile(r"[^\W_]+", re.UNICODE) +_KEYWORD_PAGE_SIZE = 256 +# Exact lexical matches outrank typical semantic distances (lower is better). +_EXACT_DISTANCE = 0.0 +_KEYWORD_DISTANCE = 0.05 + def index_capability( source: VideoSource, @@ -77,6 +84,124 @@ def speech_embedding( return encoded[0].tolist() +def _query_tokens(query: str) -> tuple[str, ...]: + return tuple( + token + for token in _TOKEN_PATTERN.findall(query.casefold()) + if len(token) > 1 + ) + + +def _keyword_distance(query: str, text: str) -> float | None: + """Return a lexical distance, or None when the text does not match.""" + + query_tokens = _query_tokens(query) + if not query_tokens: + return None + text_tokens = _query_tokens(text) + if not text_tokens: + return None + window = len(query_tokens) + for offset in range(len(text_tokens) - window + 1): + if text_tokens[offset:offset + window] == query_tokens: + return _EXACT_DISTANCE + if set(query_tokens) <= set(text_tokens): + return _KEYWORD_DISTANCE + return None + + +def _iter_speech_metadata( + storage: IndexStore, + *, + video_id: str | None, + filters: Mapping[str, Any] | None, +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + offset = 0 + while True: + batch = storage.records( + "speech", + video_id=video_id, + filters=filters, + limit=_KEYWORD_PAGE_SIZE, + offset=offset, + ) + if not batch: + break + rows.extend(batch) + if len(batch) < _KEYWORD_PAGE_SIZE: + break + offset += len(batch) + return rows + + +def keyword_search_rows( + query: str, + *, + storage: IndexStore, + video_id: str | None = None, + filters: Mapping[str, Any] | None = None, + top_k: int, +) -> list[dict[str, Any]]: + matches: list[dict[str, Any]] = [] + for metadata in _iter_speech_metadata( + storage, + video_id=video_id, + filters=filters, + ): + text = metadata.get("text") + if not isinstance(text, str): + continue + distance = _keyword_distance(query, text) + if distance is None: + continue + source_id = str(metadata.get("source_id") or "") + if not source_id: + continue + annotated = dict(metadata) + annotated["match_kind"] = ( + "exact" if distance == _EXACT_DISTANCE else "keyword" + ) + matches.append( + { + "source_id": source_id, + "metadata": annotated, + "raw_distance": distance, + } + ) + matches.sort(key=lambda row: (row["raw_distance"], row["source_id"])) + return matches[:top_k] + + +def _merge_search_rows( + *groups: list[dict[str, Any]], + top_k: int, +) -> list[dict[str, Any]]: + best: dict[str, dict[str, Any]] = {} + preferred = {"exact", "keyword"} + for group in groups: + for row in group: + source_id = str(row["source_id"]) + current = best.get(source_id) + if current is None or float(row["raw_distance"]) < float( + current["raw_distance"] + ): + best[source_id] = row + continue + if float(row["raw_distance"]) != float(current["raw_distance"]): + continue + if ( + row["metadata"].get("match_kind") in preferred + and current["metadata"].get("match_kind") not in preferred + ): + best[source_id] = row + ordered = sorted( + best.values(), + key=lambda row: (row["raw_distance"], row["source_id"]), + ) + return ordered[:top_k] + + def search_speech( query: str, *, @@ -93,17 +218,35 @@ def search_speech( raise ValueError("Search query must not be empty.") if top_k <= 0: raise ValueError("top_k must be greater than zero.") - return search_embeddings( - cleaned, + if "speech" not in config.enabled_modalities: + raise ValueError( + "The speech modality is not present in this index run." + ) + semantic_rows = storage.query( "speech", speech_embedding(cleaned, config, runtime), - config=config, - required_metadata=REQUIRED_METADATA, top_k=top_k, video_id=video_id, - query_id=query_id, filters=filters, + ) + for row in semantic_rows: + metadata = dict(row["metadata"]) + metadata.setdefault("match_kind", "semantic") + row["metadata"] = metadata + + lexical_rows = keyword_search_rows( + cleaned, storage=storage, + video_id=video_id, + filters=filters, + top_k=top_k, + ) + merged = _merge_search_rows(semantic_rows, lexical_rows, top_k=top_k) + return SearchResult( + query_id=query_id or stable_query_id(cleaned, "speech", config), + query=cleaned, + modality="speech", + hits=hits_from_rows("speech", merged, REQUIRED_METADATA), ) diff --git a/src/vidxp/capabilities/speech/segmentation.py b/src/vidxp/capabilities/speech/segmentation.py new file mode 100644 index 00000000..f19d28d0 --- /dev/null +++ b/src/vidxp/capabilities/speech/segmentation.py @@ -0,0 +1,130 @@ +"""Build searchable dialogue segments from a timed word transcript.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import Sequence + +from vidxp.capabilities.speech.config import SegmentationMode +from vidxp.capabilities.speech.transcript import TimedWord + +# Punctuation that ends a spoken sentence when attached to a word. +_SENTENCE_END = re.compile(r'[.!?…]["\')\]]*$') + + +@dataclass(frozen=True) +class DialoguePhrase: + phrase_id: int + text: str + start: float + end: float + word_start: int + word_end: int + segmentation_mode: SegmentationMode + + @property + def local_id(self) -> str: + """Stable ID for the same transcript span and segmentation mode.""" + + return ( + f"{self.segmentation_mode}:" + f"w{self.word_start:08d}-{self.word_end:08d}" + ) + + +def _phrase_from_words( + words: Sequence[TimedWord], + *, + mode: SegmentationMode, + phrase_id: int, +) -> DialoguePhrase: + if not words: + raise ValueError("A dialogue phrase requires at least one word.") + return DialoguePhrase( + phrase_id=phrase_id, + text=" ".join(word.text for word in words), + start=words[0].start, + end=words[-1].end, + word_start=words[0].index, + word_end=words[-1].index, + segmentation_mode=mode, + ) + + +def build_dialogue_phrases_from_words( + words: Sequence[TimedWord], + *, + segmentation_mode: SegmentationMode = "fixed_words", + words_per_phrase: int = 5, + window_stride_words: int = 2, +) -> list[DialoguePhrase]: + """Segment timed words into searchable phrases. + + ``fixed_words`` is the historical five-word baseline. Other modes exist so + retrieval quality can be compared (see issues #76 and #89). + """ + + if words_per_phrase <= 0: + raise ValueError("words_per_phrase must be greater than zero.") + if window_stride_words <= 0: + raise ValueError("window_stride_words must be greater than zero.") + if not words: + return [] + + phrases: list[DialoguePhrase] = [] + if segmentation_mode == "fixed_words": + for offset in range(0, len(words), words_per_phrase): + phrases.append( + _phrase_from_words( + words[offset:offset + words_per_phrase], + mode="fixed_words", + phrase_id=len(phrases), + ) + ) + return phrases + + if segmentation_mode == "overlapping_windows": + offset = 0 + while True: + phrases.append( + _phrase_from_words( + words[offset:offset + words_per_phrase], + mode="overlapping_windows", + phrase_id=len(phrases), + ) + ) + if offset + words_per_phrase >= len(words): + break + offset += window_stride_words + if offset >= len(words): + break + return phrases + + if segmentation_mode == "sentence": + current: list[TimedWord] = [] + for word in words: + current.append(word) + if _SENTENCE_END.search(word.text) or len(current) >= words_per_phrase: + phrases.append( + _phrase_from_words( + current, + mode="sentence", + phrase_id=len(phrases), + ) + ) + current = [] + if current: + phrases.append( + _phrase_from_words( + current, + mode="sentence", + phrase_id=len(phrases), + ) + ) + return phrases + + raise ValueError( + "segmentation_mode must be one of: fixed_words, " + "overlapping_windows, sentence." + ) diff --git a/src/vidxp/capabilities/speech/transcript.py b/src/vidxp/capabilities/speech/transcript.py new file mode 100644 index 00000000..74999c71 --- /dev/null +++ b/src/vidxp/capabilities/speech/transcript.py @@ -0,0 +1,151 @@ +"""Timed word transcript: the shared source for dialogue segments.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from vidxp.core.manifest import write_json_atomic + + +TRANSCRIPT_SCHEMA_VERSION = 1 +TRANSCRIPT_FILE = "speech_transcript.json" + + +@dataclass(frozen=True) +class TimedWord: + """One recognized word with its video time range. + + Speaker labels can attach later without changing text or timing (#86). + """ + + text: str + start: float + end: float + index: int + + +def _valid_interval(start: Any, end: Any, label: str) -> tuple[float, float]: + start_value = float(start) + end_value = float(end) + if start_value < 0 or end_value <= start_value: + raise ValueError( + f"{label} must have a non-negative, non-zero interval; " + f"received [{start_value}, {end_value}]." + ) + return start_value, end_value + + +def _word_text(word: Mapping[str, Any]) -> str: + return str(word.get("word", word.get("text", ""))).strip() + + +def flatten_transcript_words( + segments: Sequence[Mapping[str, Any]], +) -> list[TimedWord]: + """Normalize Whisper (or supplied) segments into ordered timed words.""" + + words: list[TimedWord] = [] + for segment_index, segment in enumerate(segments): + raw_words = segment.get("words") or [] + timestamped = [ + word + for word in raw_words + if _word_text(word) + and word.get("start") is not None + and word.get("end") is not None + ] + if timestamped: + for word in timestamped: + start, end = _valid_interval( + word["start"], + word["end"], + f"Transcript word {len(words)} in segment {segment_index}", + ) + words.append( + TimedWord( + text=_word_text(word), + start=start, + end=end, + index=len(words), + ) + ) + continue + + text = str(segment.get("text", "")).strip() + if not text: + continue + if segment.get("start") is None or segment.get("end") is None: + raise ValueError( + f"Transcript segment {segment_index} lacks start/end timestamps." + ) + start, end = _valid_interval( + segment["start"], + segment["end"], + f"Transcript segment {segment_index}", + ) + tokens = text.split() + if not tokens: + continue + duration = end - start + for offset, token in enumerate(tokens): + words.append( + TimedWord( + text=token, + start=start + duration * offset / len(tokens), + end=start + duration * (offset + 1) / len(tokens), + index=len(words), + ) + ) + return words + + +def save_transcript( + run_directory: str | Path, + words: Sequence[TimedWord], + *, + language: str | None = None, +) -> Path: + path = Path(run_directory) / TRANSCRIPT_FILE + write_json_atomic( + path, + { + "schema_version": TRANSCRIPT_SCHEMA_VERSION, + "language": language, + "words": [ + { + "index": word.index, + "word": word.text, + "start": word.start, + "end": word.end, + } + for word in words + ], + }, + ) + return path + + +def load_transcript(run_directory: str | Path) -> list[TimedWord]: + path = Path(run_directory) / TRANSCRIPT_FILE + payload = json.loads(path.read_text(encoding="utf-8")) + if int(payload.get("schema_version", 0)) != TRANSCRIPT_SCHEMA_VERSION: + raise ValueError(f"Unsupported speech transcript schema in {path}.") + words = [] + for item in payload.get("words") or (): + start, end = _valid_interval( + item["start"], + item["end"], + f"Stored transcript word {item.get('index', len(words))}", + ) + words.append( + TimedWord( + text=str(item["word"]).strip(), + start=start, + end=end, + index=int(item.get("index", len(words))), + ) + ) + return words diff --git a/tests/test_indexing.py b/tests/test_indexing.py index b60b2fb4..ccc1b0a9 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1,5 +1,7 @@ import unittest from contextlib import nullcontext +from pathlib import Path +from tempfile import TemporaryDirectory import sys from types import SimpleNamespace from unittest.mock import Mock, patch @@ -21,6 +23,11 @@ index_speech, transcribe_video, ) +from vidxp.capabilities.speech.transcript import ( + TRANSCRIPT_FILE, + flatten_transcript_words, + load_transcript, +) from vidxp.capabilities.scene.indexing import ( encode_scene_batch, scene_records, @@ -93,6 +100,179 @@ def test_word_timestamps_are_grouped_without_interpolation(self): [(item.text, item.start, item.end) for item in phrases], [("one two", 0.1, 1.0), ("three", 1.2, 1.8)], ) + self.assertEqual( + [item.local_id for item in phrases], + [ + "fixed_words:w00000000-00000001", + "fixed_words:w00000002-00000002", + ], + ) + + def test_overlapping_windows_keep_shared_words(self): + phrases = build_dialogue_phrases( + [ + { + "text": "one two three four five six", + "start": 0.0, + "end": 6.0, + "words": [ + {"word": "one", "start": 0.0, "end": 0.5}, + {"word": "two", "start": 0.5, "end": 1.0}, + {"word": "three", "start": 1.0, "end": 1.5}, + {"word": "four", "start": 1.5, "end": 2.0}, + {"word": "five", "start": 2.0, "end": 2.5}, + {"word": "six", "start": 2.5, "end": 3.0}, + ], + } + ], + words_per_phrase=4, + segmentation_mode="overlapping_windows", + window_stride_words=2, + ) + + self.assertEqual( + [item.text for item in phrases], + [ + "one two three four", + "three four five six", + ], + ) + self.assertEqual( + phrases[0].local_id, + "overlapping_windows:w00000000-00000003", + ) + + def test_sentence_boundaries_prefer_punctuation(self): + phrases = build_dialogue_phrases( + [ + { + "text": "Hello there. Next line continues", + "start": 0.0, + "end": 4.0, + "words": [ + {"word": "Hello", "start": 0.0, "end": 0.4}, + {"word": "there.", "start": 0.4, "end": 0.9}, + {"word": "Next", "start": 1.0, "end": 1.3}, + {"word": "line", "start": 1.3, "end": 1.6}, + {"word": "continues", "start": 1.6, "end": 2.2}, + ], + } + ], + words_per_phrase=8, + segmentation_mode="sentence", + ) + + self.assertEqual( + [(item.text, item.start, item.end) for item in phrases], + [ + ("Hello there.", 0.0, 0.9), + ("Next line continues", 1.0, 2.2), + ], + ) + + def test_identical_transcript_and_settings_reuse_segment_ids(self): + segments = [ + { + "text": "alpha beta gamma", + "start": 0.0, + "end": 3.0, + "words": [ + {"word": "alpha", "start": 0.0, "end": 1.0}, + {"word": "beta", "start": 1.0, "end": 2.0}, + {"word": "gamma", "start": 2.0, "end": 3.0}, + ], + } + ] + first = build_dialogue_phrases(segments, words_per_phrase=2) + second = build_dialogue_phrases(segments, words_per_phrase=2) + + self.assertEqual( + [item.local_id for item in first], + [item.local_id for item in second], + ) + + def test_segmentation_settings_change_index_fingerprint(self): + baseline = IndexConfig( + video_id="video-1", + enabled_modalities=("speech",), + capability_options={ + "speech": {"segmentation_mode": "fixed_words"}, + }, + ) + overlapping = IndexConfig( + video_id="video-1", + enabled_modalities=("speech",), + capability_options={ + "speech": {"segmentation_mode": "overlapping_windows"}, + }, + ) + + self.assertNotEqual(baseline.fingerprint(), overlapping.fingerprint()) + + def test_transcript_indexing_persists_timed_words(self): + with TemporaryDirectory() as directory: + config = IndexConfig( + dataset="hirest", + split="test", + run_id="asr", + video_id="video-1", + enabled_modalities=("speech",), + storage_directory=directory, + generation_directory=directory, + capability_options={ + "speech": {"embedding_batch_size": 2}, + }, + ) + source = VideoSource( + video_id="video-1", + transcript=( + { + "text": "first second", + "start": 0.0, + "end": 2.0, + "words": [ + {"word": "first", "start": 0.0, "end": 1.0}, + {"word": "second", "start": 1.0, "end": 2.0}, + ], + }, + ), + ) + storage = CapturingStorage() + encoder = FakeEncoder() + with ( + patch( + "vidxp.capabilities.speech.indexing.get_embedder", + return_value=encoder, + ), + patch( + "vidxp.capabilities.speech.indexing.transcribe_video", + side_effect=AssertionError("transcription was used"), + ), + ): + stats = index_speech( + source, + config=config, + storage=storage, + cancellation=CancellationToken(), + runtime=self.runtime(), + ) + + words = load_transcript(directory) + self.assertEqual( + [(word.text, word.start, word.end) for word in words], + [("first", 0.0, 1.0), ("second", 1.0, 2.0)], + ) + self.assertTrue((Path(directory) / TRANSCRIPT_FILE).is_file()) + self.assertEqual(stats["transcript_words"], 2) + self.assertEqual(stats["dialogue_phrases"], 1) + self.assertEqual( + storage.calls[0][1][0].metadata["segmentation_mode"], + "fixed_words", + ) + self.assertEqual( + flatten_transcript_words(source.transcript), + words, + ) def test_siglip2_uses_transformers_five_pooled_image_output(self): import torch @@ -174,8 +354,9 @@ def test_transcript_indexing_batches_without_transcription(self): runtime=self.runtime(), ) - self.assertEqual([len(batch) for batch in encoder.batches], [2, 1]) - self.assertEqual(stats["dialogue_phrases"], 3) + self.assertEqual([len(batch) for batch in encoder.batches], [1]) + self.assertEqual(stats["dialogue_phrases"], 1) + self.assertEqual(stats["transcript_words"], 3) def test_silent_video_skips_dialogue_before_loading_whisper(self): events = [] diff --git a/tests/test_search.py b/tests/test_search.py index accb7d84..ef7cd441 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -29,9 +29,18 @@ def __init__(self, rows): self.calls = [] def query(self, modality, embedding, **options): - self.calls.append((modality, embedding, options)) + self.calls.append(("query", modality, embedding, options)) return list(self.rows) + def records(self, modality, **options): + self.calls.append(("records", modality, options)) + video_id = options.get("video_id") + return [ + dict(row["metadata"]) + for row in self.rows + if video_id is None or row["metadata"]["video_id"] == video_id + ] + def dialogue_row(source_id, distance, video_id=MEDIA_ID): return { @@ -99,14 +108,64 @@ def test_top_k_filter_order_distance_and_score_are_preserved(self): ], ) self.assertEqual([hit.rank for hit in result.hits], [1, 2, 3]) - self.assertEqual(result.hits[0].raw_distance, 0.1) - self.assertEqual(result.hits[0].score, -0.1) + self.assertEqual(result.hits[0].raw_distance, 0.0) + self.assertEqual(result.hits[0].score, 0.0) self.assertEqual( result.hits[0].metadata, - {"text": "fresh bread", "phrase_id": 3}, + { + "text": "fresh bread", + "phrase_id": 3, + "match_kind": "exact", + }, ) - self.assertEqual(storage.calls[0][2]["top_k"], 3) - self.assertEqual(storage.calls[0][2]["video_id"], MEDIA_ID) + self.assertEqual(storage.calls[0][0], "query") + self.assertEqual(storage.calls[0][3]["top_k"], 3) + self.assertEqual(storage.calls[0][3]["video_id"], MEDIA_ID) + + def test_keyword_match_surfaces_missed_semantic_candidate(self): + semantic_only = dialogue_row("run:video-1:speech:semantic", 0.4) + semantic_only["metadata"]["text"] = "something related" + lexical = dialogue_row("run:video-1:speech:exact", 0.9) + lexical["metadata"]["text"] = "please pass the fresh bread now" + storage = FakeStorage([semantic_only, lexical]) + with patch( + "vidxp.capabilities.speech.operations.speech_embedding", + return_value=[0.5, 0.25], + ): + result = search_speech( + "fresh bread", + config=self.config, + runtime=self.runtime, + top_k=2, + video_id=MEDIA_ID, + storage=storage, + ) + + self.assertEqual(result.hits[0].source_id, "run:video-1:speech:exact") + self.assertEqual(result.hits[0].metadata["match_kind"], "exact") + self.assertEqual(result.hits[0].metadata["text"], lexical["metadata"]["text"]) + self.assertEqual(result.hits[0].start, 1.0) + self.assertEqual(result.hits[0].end, 2.0) + + def test_keyword_tokens_do_not_match_substrings(self): + row = dialogue_row("run:video-1:speech:start", 0.9) + row["metadata"]["text"] = "please start the oven" + storage = FakeStorage([row]) + with patch( + "vidxp.capabilities.speech.operations.speech_embedding", + return_value=[0.5], + ): + result = search_speech( + "art", + config=self.config, + runtime=self.runtime, + top_k=1, + video_id=MEDIA_ID, + storage=storage, + ) + + self.assertEqual(result.hits[0].metadata["match_kind"], "semantic") + self.assertEqual(result.hits[0].raw_distance, 0.9) def test_score_is_strictly_monotonic_and_not_a_probability(self): self.assertGreater(distance_to_score(0.1), distance_to_score(0.2))