scripts/test_nc_variants.py

"""Stub tests for the CC-BY-NC best-in-class variants (CPU-safe, no real models):
  1. config: nemotron 4096 branch + license decision recorded
  2. get_reranker dispatches jina-reranker-v3.5 to _JinaListwiseReranker
  3. predict(): multi-query grouping, index mapping, listwise call count
  4. visual_lane: _IS_NEMOTRON flag + nemotron loader branch present
"""
import sys
import types
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

torch_stub = types.ModuleType("torch")
torch_stub.float16 = "f16"
torch_stub.float32 = "f32"
torch_stub.bfloat16 = "bf16"
torch_stub.no_grad = lambda: (lambda f: f)
sys.modules.setdefault("torch", torch_stub)

import config

code_cfg = (ROOT / "config.py").read_text()
assert config.VISUAL_EMBED_DIM == 128, config.VISUAL_EMBED_DIM
assert "4096 if" in code_cfg, "config missing nemotron 4096 branch"
assert "ACCEPTED" in code_cfg, "license decision not recorded in config"
print("test1 dim selection: PASS")

class _FakeRankingModel:
    def __init__(self):
        self.device = "cpu"
        self.calls = []

    def eval(self):
        return self

    def to(self, device):
        self.device = device
        return self

    def rerank(self, q, docs, top_n=None):
        self.calls.append((q, tuple(docs)))
        return [
            {"index": i, "relevance_score": 1.0 if d == "B" else 0.0, "document": d}
            for i, d in enumerate(docs)
        ]

class _FakeAutoModel:
    @staticmethod
    def from_pretrained(name, **kw):
        return _FakeRankingModel()

fake_tf = types.ModuleType("transformers")
fake_tf.AutoModel = _FakeAutoModel
sys.modules["transformers"] = fake_tf

import rerankers

r = rerankers.get_reranker("jinaai/jina-reranker-v3.5")
assert type(r) is rerankers._JinaListwiseReranker
print("test2 dispatch: PASS")

pairs = [("q1", "A"), ("q1", "B"), ("q2", "C"), ("q1", "D"), ("q2", "E")]
scores = r.predict(pairs)
assert len(scores) == 5
assert scores[1] == 1.0 and scores[0] == 0.0 and scores[3] == 0.0
assert len(r.model.calls) == 2, r.model.calls
assert r.model.calls[0][1] == ("A", "B", "D")
assert r.model.calls[1][1] == ("C", "E")
print("test3 listwise grouping/mapping: PASS")

vl_src = (ROOT / "visual_lane.py").read_text()
assert '_IS_NEMOTRON = "nemotron" in _MODEL_NAME.lower()' in vl_src
assert "forward_images" in vl_src and "forward_queries" in vl_src
assert "elif _IS_TOMORO:" in vl_src, "tomoro branch must become elif"
import visual_lane
assert visual_lane._IS_NEMOTRON is False
print("test4 visual_lane nemotron branch: PASS")

print("ALL NC-VARIANT TESTS PASS")