Spaces:
Configuration error
Configuration error
File size: 5,962 Bytes
c15ad20 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | """Stub tests for the Sept-2026 upgrade options:
1. env-overridable model picks (RERANK_MODEL / VISUAL_EMBED_MODEL)
2. visual_rerank.rerank_keys ordering (fail-open, png gating)
3. dim-aware pages collection params
4. zerank id-alias fallback in the reranker factory"""
import importlib
import os
import sys
import tempfile
from pathlib import Path
REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))
os.environ["FACTORY_RAG_DATA"] = tempfile.mkdtemp(prefix="fr-upg-")
import config # noqa: E402
import ingest_lib # noqa: E402
import visual_rerank # noqa: E402
def main():
failures = []
# 1. env overrides + swap tiers
os.environ["RERANK_MODEL"] = "zeroentropy/zerank-2-reranker"
os.environ["VISUAL_EMBED_MODEL"] = "TomoroAI/tomoro-ai-colqwen3-embed-8b-awq"
cfg2 = importlib.reload(config)
if cfg2.RERANK_MODEL != "zeroentropy/zerank-2-reranker":
failures.append("RERANK_MODEL env override broken")
if cfg2.VISUAL_EMBED_DIM != 320 or cfg2.VISUAL_EMBED_MODEL != "TomoroAI/tomoro-ai-colqwen3-embed-8b-awq":
failures.append("VISUAL_EMBED_MODEL/DIM swap broken")
if "visual_embed_model" not in cfg2.snapshot():
failures.append("snapshot must expose visual_embed_model (advisor)")
os.environ["ENABLE_VISUAL_RERANK"] = "1"
cfg2 = importlib.reload(config)
if cfg2.ENABLE_VISUAL_RERANK is not True:
failures.append("ENABLE_VISUAL_RERANK env override broken")
del os.environ["RERANK_MODEL"], os.environ["VISUAL_EMBED_MODEL"], os.environ["ENABLE_VISUAL_RERANK"]
importlib.reload(config)
# 2. visual_rerank.rerank_keys ordering
sha = "abc123"
png_dir = ingest_lib.PAGES_DIR / sha
png_dir.mkdir(parents=True, exist_ok=True)
from PIL import Image
for pg in (1, 2):
Image.new("RGB", (4, 4)).save(png_dir / f"p{pg:04d}.png")
class FakeRE:
def __init__(self, scores):
self.scores = scores
self.seen = None
def predict(self, pairs, prompt=None, **kw):
self.seen = [p[1]["image"] for p in pairs]
return self.scores
fake = FakeRE([0.1, 0.9]) # page 2 more relevant
vis_flag_orig = config.ENABLE_VISUAL_RERANK
config.ENABLE_VISUAL_RERANK = True
visual_rerank.get_visual_reranker = lambda: fake
try:
keys = [(sha, 1), (sha, 2), (sha, None)] # third has no page -> rest
out = visual_rerank.rerank_keys("torque?", keys)
if out[:2] != [(sha, 2), (sha, 1)]:
failures.append(f"rerank_keys order wrong: {out}")
if out[2] != (sha, None):
failures.append("no-png entry must land after scored ones")
if not any(p.endswith("p0002.png") for p in fake.seen):
failures.append("reranker did not receive page PNG path")
# fail-open: reranker exception -> original order
visual_rerank.get_visual_reranker = lambda: None
if visual_rerank.rerank_keys("q", keys) != keys:
failures.append("rerank_keys must be fail-open with no reranker")
config.ENABLE_VISUAL_RERANK = False
if visual_rerank.rerank_keys("q", keys) != keys:
failures.append("rerank_keys must no-op when flag off")
finally:
visual_rerank.get_visual_reranker = None
config.ENABLE_VISUAL_RERANK = vis_flag_orig if False else vis_flag_orig
config.ENABLE_VISUAL_RERANK = vis_flag_orig
# 3. dim-aware pages collection params
import types
_tm = types.ModuleType("torch")
_tm.float32 = "f32"; _tm.float16 = "f16"; _tm.bfloat16 = "bf16"
sys.modules.setdefault("torch", _tm) # CPU sandbox: no torch; stubbed attrs suffice
import visual_lane
dim_orig = config.VISUAL_EMBED_DIM
try:
config.VISUAL_EMBED_DIM = 320
params = visual_lane.config_pages_vector_params()
if params.size != 320:
failures.append("tomoro dim (320) not used for pages collection")
config.VISUAL_EMBED_DIM = 128
params = visual_lane.config_pages_vector_params()
if params.size != 128:
failures.append("colqwen2.5 dim (128) not used")
finally:
config.VISUAL_EMBED_DIM = dim_orig
# 4. zerank id-alias fallback (fake transformers in sys.modules)
saved_tf = sys.modules.get("transformers")
# 4. zerank id-alias fallback (fake transformers module in sys.modules)
import types
import rerankers
class FakeTok:
def __init__(self):
self.name = None
@classmethod
def from_pretrained(cls, name, padding_side=None):
if name == "zeroentropy/zerank-2-reranker":
raise OSError("alias missing")
obj = cls()
obj.name = name
return obj
def encode(self, s, add_special_tokens=False):
return [1]
class FakeCausalLM:
@classmethod
def from_pretrained(cls, name, **kw):
obj = cls()
return obj
def to(self, device):
return self
def eval(self):
return self
fake_tf = types.ModuleType("transformers")
fake_tf.AutoTokenizer = FakeTok
fake_tf.AutoModelForCausalLM = FakeCausalLM
saved_tf = sys.modules.get("transformers")
sys.modules["transformers"] = fake_tf
try:
r = rerankers._QwenPairwiseReranker("zeroentropy/zerank-2-reranker")
if r.tok is None or r.tok.name != "zeroentropy/zerank-2":
failures.append("zerank id-alias fallback broken")
finally:
if saved_tf is not None:
sys.modules["transformers"] = saved_tf
else:
sys.modules.pop("transformers", None)
if failures:
print("FAIL:")
for f in failures:
print(" -", f)
sys.exit(1)
print("PASS: env overrides + swap tiers, rerank_keys ordering + fail-open, "
"dim-aware collection, zerank alias fallback")
if __name__ == "__main__":
main()
|