diff --git a/README.md b/README.md index 7c4dd42..8c63fda 100644 --- a/README.md +++ b/README.md @@ -154,19 +154,19 @@ No confident match found (best score: 0.28). Show results anyway? [y/N]: With `--no-trim`, low-confidence results are shown with a note instead of a prompt. -Options: `--results N`, `--output-dir DIR`, `--no-trim` to skip auto-trimming, `--threshold 0.5` to adjust the confidence cutoff, `--save-top N` to save the top N clips instead of just the best match, `--dedupe` to drop results too similar to a higher-ranked pick (prevents near-duplicate chunks of the same event from filling the list), and `--rerank` to ask Gemini Flash to re-rank the returned candidates before trimming. Backend and model are auto-detected from the index — pass `--backend` or `--model` only to override. +Options: `--results N`, `--output-dir DIR`, `--no-trim` to skip auto-trimming, `--threshold 0.5` to adjust the confidence cutoff, `--save-top N` to save the top N clips instead of just the best match, `--dedupe` to drop results too similar to a higher-ranked pick (prevents near-duplicate chunks of the same event from filling the list), and `--rerank` to ask a VLM to re-rank the returned candidates before trimming. Backend and model are auto-detected from the index — pass `--backend` or `--model` only to override. ```bash # Save top 5 clips, dropping near-duplicates sentrysearch search "red truck" --save-top 5 --dedupe 0.9 -# Re-rank the top 10 embedding matches with Gemini Flash before trimming +# Re-rank the top 10 embedding matches with a VLM before trimming sentrysearch search "pedestrian crossing behind the car" --rerank --results 10 ``` The `--dedupe` value is a cosine similarity ceiling (0–1). Any result whose similarity to an already-kept higher-ranked result exceeds this value is dropped. Lower values are stricter: `0.8` requires results to be very distinct, `0.95` only removes near-identical chunks. `0.9` is a good default. -`--rerank` extracts each returned candidate clip, sends it to Gemini 2.5 Flash with the query, and sorts likely visual matches ahead of embedding-only results. If reranking cannot run or a candidate cannot be scored, SentrySearch keeps the embedding-ranked results instead of failing the search. +`--rerank` extracts each returned candidate clip, sends it to a VLM with the query, and sorts likely visual matches ahead of embedding-only results. Gemini and qwen-cloud searches use Gemini 2.5 Flash for reranking; local searches use a local Qwen3-VL Instruct reranker. If reranking cannot run or a candidate cannot be scored, SentrySearch keeps the embedding-ranked results instead of failing the search. ### Search by image @@ -279,6 +279,7 @@ Options: Notes: - First run downloads the model (~16 GB for 8B, ~4 GB for 2B). +- Local `--rerank` downloads a separate Qwen3-VL Instruct model (`Qwen/Qwen3-VL-8B-Instruct` or `Qwen/Qwen3-VL-2B-Instruct`) in addition to the embedding model. - Embeddings from different backends and models are **not compatible**. Each backend/model combination gets its own isolated index, so they can't accidentally mix. If you search with a model that has no indexed data, you'll be told which model was actually used. - Speed varies by GPU core count — base M-series chips are slower than Pro/Max but produce identical results. diff --git a/sentrysearch/cli.py b/sentrysearch/cli.py index 13b147d..d530bff 100644 --- a/sentrysearch/cli.py +++ b/sentrysearch/cli.py @@ -98,6 +98,30 @@ def _strip_private_result_fields(results: list[dict]) -> list[dict]: ] +def _get_search_reranker( + backend: str, + model: str | None, + quantize: bool | None, +): + if backend == "local": + from .embedder import reset_embedder + from .qwen_reranker import QwenReranker + + reset_embedder() + return QwenReranker( + model_name=model or "qwen8b", + quantize=quantize, + ).load() + + from .gemini_reranker import GeminiReranker + return GeminiReranker() + + +def _first_error_line(exc: Exception) -> str: + lines = str(exc).splitlines() + return lines[0] if lines else exc.__class__.__name__ + + def _open_file(path: str) -> None: """Open a file with the system's default application.""" try: @@ -659,7 +683,7 @@ def index(directory, chunk_duration, overlap, preprocess, target_resolution, help="Drop results whose cosine similarity to a higher-ranked " "result exceeds this (e.g. 0.9).") @click.option("--rerank", is_flag=True, - help="Use Gemini Flash to rerank candidates before trimming.") + help="Use a VLM to rerank candidates before trimming.") @click.option("--verbose", is_flag=True, help="Show debug info.") def search(query, n_results, output_dir, trim, save_top, threshold, overlay, backend, model, dashscope_model, quantize, dedupe_threshold, rerank, verbose): """Search indexed footage with a natural language QUERY.""" @@ -740,14 +764,14 @@ def search(query, n_results, output_dir, trim, save_top, threshold, overlay, bac rerank_enabled = rerank if rerank and results: from .gemini_embedder import GeminiAPIKeyError - from .gemini_reranker import GeminiReranker + from .local_embedder import LocalModelError from .reranker import rerank_results try: - reranker = GeminiReranker() - except GeminiAPIKeyError: + reranker = _get_search_reranker(backend, model, quantize) + except (GeminiAPIKeyError, LocalModelError) as exc: click.secho( - "--rerank skipped: GEMINI_API_KEY is not set; " + f"--rerank skipped: {_first_error_line(exc)}; " "showing embedding results.", fg="yellow", err=True, diff --git a/sentrysearch/gemini_reranker.py b/sentrysearch/gemini_reranker.py index f846504..14665b6 100644 --- a/sentrysearch/gemini_reranker.py +++ b/sentrysearch/gemini_reranker.py @@ -14,37 +14,17 @@ _RateLimiter, _retry, ) -from .reranker import RerankScore, parse_rerank_response +from .reranker import ( + RERANK_SCHEMA, + RerankScore, + build_rerank_prompt, + parse_rerank_response, +) load_dotenv() RERANK_MODEL = "gemini-2.5-flash" -_RERANK_SCHEMA = { - "type": "object", - "properties": { - "rerank_match": {"type": "boolean"}, - "rerank_confidence": { - "type": "number", - "minimum": 0.0, - "maximum": 1.0, - }, - }, - "required": ["rerank_match", "rerank_confidence"], -} - - -def _prompt(query: str) -> str: - return ( - "You are reranking video search candidates.\n" - "Look at the clip and decide whether it visually matches this search " - f"query: {query!r}\n\n" - "Return JSON only with this exact shape:\n" - '{"rerank_match": true, "rerank_confidence": 0.0}\n' - "Use rerank_match=true only when the clip contains the requested event. " - "Use rerank_confidence as your confidence from 0.0 to 1.0." - ) - class GeminiReranker: """Gemini Flash reranker for candidate video clips.""" @@ -74,10 +54,10 @@ def score( from google.genai import types video_part = GeminiEmbedder._make_video_part(clip_path, types) - prompt_part = types.Part(text=_prompt(query)) + prompt_part = types.Part(text=build_rerank_prompt(query)) config = types.GenerateContentConfig( response_mime_type="application/json", - response_json_schema=_RERANK_SCHEMA, + response_json_schema=RERANK_SCHEMA, temperature=0.0, ) diff --git a/sentrysearch/qwen_reranker.py b/sentrysearch/qwen_reranker.py new file mode 100644 index 0000000..842294d --- /dev/null +++ b/sentrysearch/qwen_reranker.py @@ -0,0 +1,265 @@ +"""Local Qwen3-VL Instruct reranker.""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +from .local_embedder import LocalModelError, _cpu_fallback_warning +from .reranker import RerankScore, build_rerank_prompt, parse_rerank_response + + +RERANK_MODEL_ALIASES: dict[str, str] = { + "qwen8b": "Qwen/Qwen3-VL-8B-Instruct", + "qwen2b": "Qwen/Qwen3-VL-2B-Instruct", +} +MAX_NEW_TOKENS = 128 +QWEN3_IMAGE_PATCH_SIZE = 16 + + +class QwenReranker: + """Local Qwen3-VL Instruct reranker for candidate video clips.""" + + def __init__( + self, + model_name: str = "qwen8b", + *, + quantize: bool | None = None, + ): + if model_name not in RERANK_MODEL_ALIASES: + raise LocalModelError( + "Local rerank supports qwen2b or qwen8b. " + f"Got indexed model: {model_name}" + ) + self._model_name = RERANK_MODEL_ALIASES[model_name] + self._quantize = quantize + self._model = None + self._processor = None + self._process_vision_info = None + + def load(self): + self._load_model() + return self + + def _load_model(self) -> None: + if self._model is not None: + return + + try: + import torch + from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLForConditionalGeneration, + ) + from transformers.models.qwen3_vl.processing_qwen3_vl import ( + Qwen3VLProcessor, + ) + from qwen_vl_utils import process_vision_info + except ImportError as e: + raise LocalModelError( + f"Missing dependencies for local Qwen rerank: {e}\n\n" + "Install with: uv tool install \".[local]\"\n" + "For 4-bit quantization: uv tool install \".[local-quantized]\"" + ) from e + + try: + from huggingface_hub import try_to_load_from_cache + cached = try_to_load_from_cache(self._model_name, "config.json") + is_cached = isinstance(cached, str) and os.path.exists(cached) + except Exception: + is_cached = False + + if is_cached: + print(f"Loading {self._model_name}...", file=sys.stderr) + else: + print( + f"Downloading {self._model_name} (this only happens once)...", + file=sys.stderr, + ) + + if torch.cuda.is_available(): + device = "cuda" + dtype = torch.bfloat16 + elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + device = "mps" + dtype = torch.float16 + else: + device = "cpu" + dtype = torch.float32 + print(_cpu_fallback_warning(), file=sys.stderr) + + quantization_config = None + want_quantize = self._quantize + if want_quantize is None and device == "cuda": + props = torch.cuda.get_device_properties(0) + total_mem = getattr(props, "total_memory", None) or getattr( + props, "total_mem", 0, + ) + vram_gb = total_mem / (1024 ** 3) + needs_gb = 18 if "8B" in self._model_name else 6 + want_quantize = vram_gb < needs_gb + if want_quantize and device == "cuda": + try: + import bitsandbytes # noqa: F401 + from transformers import BitsAndBytesConfig + quantization_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch.float16, + ) + print("Using 4-bit quantization (bitsandbytes)", file=sys.stderr) + except ImportError as e: + if self._quantize is True: + raise LocalModelError( + "4-bit quantization requested but bitsandbytes is not " + "installed.\n\n" + "Install with: uv tool install \".[local-quantized]\"" + ) from e + elif want_quantize and device != "cuda": + if self._quantize is True: + raise LocalModelError( + "4-bit quantization requires CUDA (NVIDIA GPU). " + f"Current device: {device}" + ) + + try: + self._processor = Qwen3VLProcessor.from_pretrained( + self._model_name, padding_side="left", + ) + self._process_vision_info = process_vision_info + + load_kwargs = dict(trust_remote_code=True) + if device == "mps": + load_kwargs["attn_implementation"] = "eager" + os.environ.setdefault("TRANSFORMERS_DISABLE_TORCH_CHECK", "1") + if quantization_config is not None: + load_kwargs["quantization_config"] = quantization_config + else: + load_kwargs["torch_dtype"] = dtype + + self._model = Qwen3VLForConditionalGeneration.from_pretrained( + self._model_name, **load_kwargs, + ) + if quantization_config is None: + self._model = self._model.to(device) + self._model.eval() + print(f"Rerank model loaded on {device}", file=sys.stderr) + except Exception as e: + raise LocalModelError( + f"Failed to load {self._model_name}: {e}" + ) from e + + def _model_device(self): + device = getattr(self._model, "device", None) + if device is not None: + return device + try: + return next(self._model.parameters()).device + except Exception: + return None + + @staticmethod + def _trim_generated_ids(input_ids, generated_ids): + return [ + output_ids[len(input_ids_row):] + for input_ids_row, output_ids in zip(input_ids, generated_ids) + ] + + def score( + self, + query: str, + clip_path: str, + *, + verbose: bool = False, + ) -> RerankScore | None: + """Return a validated rerank score, or None for unparsable model output.""" + self._load_model() + + import torch + process_vision_info = self._process_vision_info + if process_vision_info is None: + from qwen_vl_utils import process_vision_info + + clip = Path(clip_path) + if not clip.exists(): + raise LocalModelError(f"Clip file not found: {clip}") + + t0 = time.monotonic() + conversation = [ + { + "role": "user", + "content": [ + { + "type": "video", + "video": "file://" + str(clip.resolve()), + "fps": 1.0, + "max_frames": 32, + }, + {"type": "text", "text": build_rerank_prompt(query)}, + ], + }, + ] + text = self._processor.apply_chat_template( + conversation, tokenize=False, add_generation_prompt=True, + ) + + images, video_inputs, video_kwargs = process_vision_info( + conversation, + image_patch_size=QWEN3_IMAGE_PATCH_SIZE, + return_video_metadata=True, + return_video_kwargs=True, + ) + video_kwargs = video_kwargs or {} + + if video_inputs is not None: + videos, video_metadata = zip(*video_inputs) + videos = list(videos) + video_metadata = list(video_metadata) + else: + videos, video_metadata = None, None + + inputs = self._processor( + text=[text], + images=images, + videos=videos, + video_metadata=video_metadata, + return_tensors="pt", + padding=True, + **video_kwargs, + ) + + device = self._model_device() + if device is not None: + inputs = { + key: value.to(device) if hasattr(value, "to") else value + for key, value in inputs.items() + } + + with torch.no_grad(): + generated_ids = self._model.generate( + **inputs, + do_sample=False, + max_new_tokens=MAX_NEW_TOKENS, + ) + generated_ids = self._trim_generated_ids( + inputs["input_ids"], generated_ids, + ) + output_text = self._processor.batch_decode( + generated_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + )[0] + + elapsed = time.monotonic() - t0 + score = parse_rerank_response(output_text) + if verbose: + status = "fallback" if score is None else ( + f"match={score.rerank_match}, " + f"confidence={score.rerank_confidence:.2f}" + ) + print( + f" [verbose] rerank {self._model_name}: {status}, " + f"inference_time={elapsed:.2f}s", + file=sys.stderr, + ) + return score diff --git a/sentrysearch/reranker.py b/sentrysearch/reranker.py index 79f6c3a..e1ac537 100644 --- a/sentrysearch/reranker.py +++ b/sentrysearch/reranker.py @@ -19,14 +19,55 @@ class RerankScore: rerank_confidence: float +RERANK_SCHEMA = { + "type": "object", + "properties": { + "rerank_match": {"type": "boolean"}, + "rerank_confidence": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0, + }, + }, + "required": ["rerank_match", "rerank_confidence"], +} + + +def build_rerank_prompt(query: str) -> str: + return ( + "You are reranking video search candidates.\n" + "Look at the clip and decide whether it visually matches this search " + f"query: {query!r}\n\n" + "Return JSON only with this exact shape:\n" + '{"rerank_match": true, "rerank_confidence": 0.0}\n' + "Use rerank_match=true only when the clip contains the requested event. " + "Use rerank_confidence as your confidence from 0.0 to 1.0." + ) + + +def _strip_json_code_fence(text): + if not isinstance(text, str): + return text + text = text.strip() + if not text.startswith("```"): + return text + + lines = text.splitlines() + if len(lines) < 3 or lines[-1].strip() != "```": + return text + if lines[0].strip() not in ("```", "```json", "```JSON"): + return text + return "\n".join(lines[1:-1]).strip() + + def parse_rerank_response(text: str) -> RerankScore | None: - """Parse and validate the Gemini rerank JSON response. + """Parse and validate a VLM rerank JSON response. Invalid model output returns ``None`` so callers can fall back to the candidate's embedding rank. """ try: - data = json.loads(text) + data = json.loads(_strip_json_code_fence(text)) except (json.JSONDecodeError, TypeError): return None diff --git a/tests/test_cli.py b/tests/test_cli.py index f944b06..f606175 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -835,7 +835,7 @@ def test_search_rerank_missing_key_falls_back_to_embedding_results( ] with patch("sentrysearch.store.SentryStore") as MockStore, \ patch("sentrysearch.embedder.get_embedder", return_value=MagicMock()), \ - patch("sentrysearch.store.detect_index", return_value=("local", "qwen2b")), \ + patch("sentrysearch.store.detect_index", return_value=("gemini", None)), \ patch("sentrysearch.search.search_footage", return_value=results), \ patch("sentrysearch.gemini_reranker.GeminiReranker", side_effect=GeminiAPIKeyError("GEMINI_API_KEY is not set.")) as MockReranker, \ @@ -857,6 +857,85 @@ def test_search_rerank_missing_key_falls_back_to_embedding_results( assert mock_present.call_args[0][0] == results assert mock_present.call_args.kwargs["rerank_enabled"] is False + def test_search_local_rerank_uses_qwen_and_resets_before_load( + self, runner, + ): + results = [ + {"source_file": "/a.mp4", "start_time": 0.0, "end_time": 30.0, + "similarity_score": 0.7}, + ] + events = [] + + def fake_reset(): + events.append("reset") + + with patch("sentrysearch.store.SentryStore") as MockStore, \ + patch("sentrysearch.embedder.get_embedder", return_value=MagicMock()), \ + patch("sentrysearch.embedder.reset_embedder", + side_effect=fake_reset), \ + patch("sentrysearch.store.detect_index", + return_value=("local", "qwen2b")), \ + patch("sentrysearch.search.search_footage", return_value=results), \ + patch("sentrysearch.qwen_reranker.QwenReranker") as MockQwen, \ + patch("sentrysearch.reranker.rerank_results", + return_value=results) as mock_rerank, \ + patch("sentrysearch.cli._cache_last_search"), \ + patch("sentrysearch.cli._present_results"): + inst = MagicMock() + inst.get_stats.return_value = {"total_chunks": 5} + MockStore.return_value = inst + qwen = MockQwen.return_value + + def fake_load(): + events.append("load") + return qwen + + qwen.load.side_effect = fake_load + + result = runner.invoke(cli, ["search", "test", "--rerank"]) + + assert result.exit_code == 0, result.output + MockQwen.assert_called_once_with(model_name="qwen2b", quantize=None) + assert events[:2] == ["reset", "load"] + mock_rerank.assert_called_once_with( + "test", results, qwen, + candidate_dir=ANY, + verbose=False, + ) + + def test_search_local_rerank_load_error_falls_back_to_embedding_results( + self, runner, + ): + from sentrysearch.local_embedder import LocalModelError + + results = [ + {"source_file": "/a.mp4", "start_time": 0.0, "end_time": 30.0, + "similarity_score": 0.7}, + ] + with patch("sentrysearch.store.SentryStore") as MockStore, \ + patch("sentrysearch.embedder.get_embedder", return_value=MagicMock()), \ + patch("sentrysearch.store.detect_index", + return_value=("local", "qwen2b")), \ + patch("sentrysearch.search.search_footage", return_value=results), \ + patch("sentrysearch.qwen_reranker.QwenReranker", + side_effect=LocalModelError("missing local deps")), \ + patch("sentrysearch.reranker.rerank_results") as mock_rerank, \ + patch("sentrysearch.cli._cache_last_search") as mock_cache, \ + patch("sentrysearch.cli._present_results") as mock_present: + inst = MagicMock() + inst.get_stats.return_value = {"total_chunks": 5} + MockStore.return_value = inst + + result = runner.invoke(cli, ["search", "test", "--rerank"]) + + assert result.exit_code == 0, result.output + assert "--rerank skipped: missing local deps" in result.output + mock_rerank.assert_not_called() + mock_cache.assert_called_once_with(results, query="test") + mock_present.assert_called_once() + assert mock_present.call_args[0][0] == results + assert mock_present.call_args.kwargs["rerank_enabled"] is False + def test_search_dedupe_runs_before_rerank(self, runner): deduped = [ {"source_file": "/a.mp4", "start_time": 0.0, "end_time": 30.0, diff --git a/tests/test_gemini_reranker.py b/tests/test_gemini_reranker.py index ddbdb9c..cd432fc 100644 --- a/tests/test_gemini_reranker.py +++ b/tests/test_gemini_reranker.py @@ -6,15 +6,11 @@ import pytest from sentrysearch.gemini_embedder import GeminiAPIKeyError -from sentrysearch.gemini_reranker import GeminiReranker, RERANK_MODEL, _prompt -from sentrysearch.reranker import RerankScore +from sentrysearch.gemini_reranker import GeminiReranker, RERANK_MODEL +from sentrysearch.reranker import RERANK_SCHEMA, RerankScore class TestGeminiReranker: - def test_prompt_is_video_generic(self): - prompt = _prompt("red truck") - assert "video search candidates" in prompt - def test_raises_without_api_key(self): with patch.dict(os.environ, {}, clear=True): os.environ.pop("GEMINI_API_KEY", None) @@ -28,10 +24,13 @@ def test_creates_client_with_key(self, mock_client_cls): mock_client_cls.assert_called_once_with(api_key="test-key-123") @patch("sentrysearch.gemini_reranker._retry", side_effect=lambda fn: fn()) + @patch("sentrysearch.gemini_reranker.build_rerank_prompt", + return_value="shared prompt") @patch("sentrysearch.gemini_reranker._RateLimiter") @patch("google.genai.Client") def test_score_uses_flash_json_retry_and_limiter( - self, mock_client_cls, mock_limiter_cls, mock_retry, tmp_path, + self, mock_client_cls, mock_limiter_cls, mock_prompt, mock_retry, + tmp_path, ): clip = tmp_path / "candidate.mp4" clip.write_bytes(b"fake-video") @@ -46,6 +45,7 @@ def test_score_uses_flash_json_retry_and_limiter( score = reranker.score("red truck", str(clip)) assert score == RerankScore(True, 0.91) + mock_prompt.assert_called_once_with("red truck") mock_limiter_cls.return_value.wait.assert_called_once() mock_retry.assert_called_once() @@ -54,9 +54,7 @@ def test_score_uses_flash_json_retry_and_limiter( assert call.kwargs["model"] == "gemini-2.5-flash" config = call.kwargs["config"] assert config.response_mime_type == "application/json" - assert config.response_json_schema["required"] == [ - "rerank_match", "rerank_confidence", - ] + assert config.response_json_schema == RERANK_SCHEMA @patch("sentrysearch.gemini_reranker._retry", side_effect=lambda fn: fn()) @patch("sentrysearch.gemini_reranker._RateLimiter") diff --git a/tests/test_qwen_reranker.py b/tests/test_qwen_reranker.py new file mode 100644 index 0000000..4b09706 --- /dev/null +++ b/tests/test_qwen_reranker.py @@ -0,0 +1,160 @@ +"""Tests for local Qwen3-VL reranker (mocked, no model download).""" + +import sys +import types +from unittest.mock import MagicMock, patch + +import pytest + +from sentrysearch.local_embedder import LocalModelError +from sentrysearch.qwen_reranker import ( + MAX_NEW_TOKENS, + QwenReranker, + QWEN3_IMAGE_PATCH_SIZE, + RERANK_MODEL_ALIASES, +) +from sentrysearch.reranker import RerankScore + + +class _NoGrad: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + +class TestQwenReranker: + def test_aliases_resolve_to_instruct_models(self): + reranker = QwenReranker("qwen2b") + assert reranker._model_name == "Qwen/Qwen3-VL-2B-Instruct" + assert RERANK_MODEL_ALIASES["qwen8b"] == ( + "Qwen/Qwen3-VL-8B-Instruct" + ) + + def test_rejects_custom_model_for_v1(self): + with pytest.raises(LocalModelError, match="qwen2b or qwen8b"): + QwenReranker("custom_model") + + def test_load_missing_dependencies_raises_local_model_error(self): + reranker = QwenReranker("qwen2b") + with patch.dict(sys.modules, {"torch": None}): + with pytest.raises(LocalModelError, match="Missing dependencies"): + reranker.load() + + def test_load_model_wires_processor_model_and_vision_helper(self): + reranker = QwenReranker("qwen2b") + processor = MagicMock() + model = MagicMock() + model.to.return_value = model + + processor_cls = MagicMock() + processor_cls.from_pretrained.return_value = processor + model_cls = MagicMock() + model_cls.from_pretrained.return_value = model + process_vision_info = MagicMock() + + torch_module = types.ModuleType("torch") + torch_module.float16 = "float16" + torch_module.float32 = "float32" + torch_module.bfloat16 = "bfloat16" + torch_module.cuda = MagicMock() + torch_module.cuda.is_available.return_value = False + torch_module.backends = types.SimpleNamespace( + mps=types.SimpleNamespace(is_available=lambda: False), + ) + + modeling_module = types.ModuleType( + "transformers.models.qwen3_vl.modeling_qwen3_vl", + ) + modeling_module.Qwen3VLForConditionalGeneration = model_cls + processing_module = types.ModuleType( + "transformers.models.qwen3_vl.processing_qwen3_vl", + ) + processing_module.Qwen3VLProcessor = processor_cls + qwen_utils = types.ModuleType("qwen_vl_utils") + qwen_utils.process_vision_info = process_vision_info + huggingface_hub = types.ModuleType("huggingface_hub") + huggingface_hub.try_to_load_from_cache = MagicMock(return_value=None) + + with patch.dict(sys.modules, { + "torch": torch_module, + "transformers": types.ModuleType("transformers"), + "transformers.models": types.ModuleType("transformers.models"), + "transformers.models.qwen3_vl": types.ModuleType( + "transformers.models.qwen3_vl", + ), + "transformers.models.qwen3_vl.modeling_qwen3_vl": + modeling_module, + "transformers.models.qwen3_vl.processing_qwen3_vl": + processing_module, + "qwen_vl_utils": qwen_utils, + "huggingface_hub": huggingface_hub, + }): + assert reranker.load() is reranker + + processor_cls.from_pretrained.assert_called_once_with( + "Qwen/Qwen3-VL-2B-Instruct", padding_side="left", + ) + model_cls.from_pretrained.assert_called_once_with( + "Qwen/Qwen3-VL-2B-Instruct", + trust_remote_code=True, + torch_dtype=torch_module.float32, + ) + model.to.assert_called_once_with("cpu") + model.eval.assert_called_once_with() + assert reranker._processor is processor + assert reranker._model is model + assert reranker._process_vision_info is process_vision_info + + def test_score_generates_decodes_and_parses(self, tmp_path): + clip = tmp_path / "candidate.mp4" + clip.write_bytes(b"fake-video") + + reranker = QwenReranker("qwen2b") + model = MagicMock() + model.device = "cpu" + model.generate.return_value = [[10, 11, 12, 13]] + processor = MagicMock() + processor.apply_chat_template.return_value = "chat prompt" + processor.return_value = { + "input_ids": [[10, 11]], + "attention_mask": [[1, 1]], + } + processor.batch_decode.return_value = [ + '```json\n{"rerank_match": true, "rerank_confidence": 0.82}\n```', + ] + reranker._model = model + reranker._processor = processor + + torch_module = types.ModuleType("torch") + torch_module.no_grad = lambda: _NoGrad() + + qwen_utils = types.ModuleType("qwen_vl_utils") + qwen_utils.process_vision_info = MagicMock( + return_value=(None, [("video-data", {"fps": 1.0})], {"fps": 1.0}), + ) + + with patch.dict( + sys.modules, + {"torch": torch_module, "qwen_vl_utils": qwen_utils}, + ): + score = reranker.score("red truck", str(clip), verbose=True) + + assert score == RerankScore(True, 0.82) + qwen_utils.process_vision_info.assert_called_once() + assert qwen_utils.process_vision_info.call_args.kwargs == { + "image_patch_size": QWEN3_IMAGE_PATCH_SIZE, + "return_video_metadata": True, + "return_video_kwargs": True, + } + processor.apply_chat_template.assert_called_once() + processor.assert_called_once() + model.generate.assert_called_once() + assert model.generate.call_args.kwargs["do_sample"] is False + assert model.generate.call_args.kwargs["max_new_tokens"] == MAX_NEW_TOKENS + processor.batch_decode.assert_called_once_with( + [[12, 13]], + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) diff --git a/tests/test_reranker.py b/tests/test_reranker.py index c1364c9..154528e 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -1,7 +1,9 @@ """Tests for VLM reranking helpers.""" from sentrysearch.reranker import ( + RERANK_SCHEMA, RerankScore, + build_rerank_prompt, parse_rerank_response, rerank_results, ) @@ -17,15 +19,34 @@ def _result(name: str, score: float = 0.8) -> dict: class TestParseRerankResponse: + def test_prompt_is_video_generic(self): + prompt = build_rerank_prompt("red truck") + assert "video search candidates" in prompt + assert "red truck" in prompt + + def test_schema_requires_score_fields(self): + assert RERANK_SCHEMA["required"] == [ + "rerank_match", "rerank_confidence", + ] + def test_accepts_valid_json(self): score = parse_rerank_response( '{"rerank_match": true, "rerank_confidence": 0.75}', ) assert score == RerankScore(True, 0.75) + def test_accepts_json_code_fence(self): + score = parse_rerank_response( + '```json\n{"rerank_match": true, "rerank_confidence": 0.75}\n```', + ) + assert score == RerankScore(True, 0.75) + def test_rejects_malformed_json(self): assert parse_rerank_response("{not json") is None + def test_rejects_missing_text(self): + assert parse_rerank_response(None) is None + def test_rejects_non_object_json(self): assert parse_rerank_response("[true, 0.8]") is None