File size: 4,436 Bytes
0c19da8 4d6878e 0c19da8 4d6878e 901bf8c 0c19da8 | 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 | """Unit tests for the embedding-text builder (canlex/embed.py)."""
import unittest
from canlex.embed import embed_text, embed_texts, content_hash, plan
def _chunk(**kw):
base = {"doc_type": "legislation", "act_short": "X", "marginal_note": "",
"part": "", "heading": "", "text": ""}
base.update(kw)
return base
class EmbedTextTests(unittest.TestCase):
def test_memorandum_title_comes_from_part(self):
# A memo's marginal note is generic; its subject lives in 'part'.
out = embed_text(_chunk(doc_type="memorandum", act_short="D-Memo",
marginal_note="Guidelines",
part="Value for Duty", text="body"))
self.assertIn("Value for Duty", out)
def test_legislation_title_comes_from_marginal_note(self):
out = embed_text(_chunk(marginal_note="Application for protection",
part="PART 2", text="body"))
self.assertIn("Application for protection", out)
def test_title_is_repeated_for_emphasis(self):
out = embed_text(_chunk(marginal_note="UNIQUEWORD", text="b"))
self.assertEqual(out.count("UNIQUEWORD"), 2)
class IncrementalPlanTests(unittest.TestCase):
"""plan() decides which chunks the incremental build may reuse."""
def _stored(self, chunk, rows=None):
texts = embed_texts(chunk)
return {chunk["id"]: (content_hash(texts),
rows if rows is not None else [[0.0]] * len(texts))}
def test_unchanged_chunk_is_reused(self):
c = _chunk(id="a", marginal_note="Title", text="body")
reusable, to_embed = plan([c], self._stored(c))
self.assertEqual([r[0] for r in reusable], ["a"])
self.assertEqual(to_embed, [])
def test_changed_text_is_re_embedded(self):
old = _chunk(id="a", marginal_note="Title", text="old body")
new = _chunk(id="a", marginal_note="Title", text="new body")
reusable, to_embed = plan([new], self._stored(old))
self.assertEqual(reusable, [])
self.assertEqual([r[0] for r in to_embed], ["a"])
def test_new_chunk_and_window_count_mismatch_are_re_embedded(self):
c = _chunk(id="a", marginal_note="Title", text="body")
reusable, to_embed = plan([c], {})
self.assertEqual([r[0] for r in to_embed], ["a"])
# same hash but a different number of stored rows: not trusted
stored = self._stored(c, rows=[[0.0], [0.0]])
reusable, to_embed = plan([c], stored)
self.assertEqual([r[0] for r in to_embed], ["a"])
def test_missing_hashes_means_nothing_reused(self):
from canlex.embed import load_stored
import pathlib
self.assertEqual(load_stored(pathlib.Path("does-not-exist.npz")), {})
# ---- Audit follow-up 2026-09-21, group A items 1 and 2 ---------------------
def _embedder_cached():
from huggingface_hub import try_to_load_from_cache
from canlex.embed import EMB_REPO
return all(isinstance(try_to_load_from_cache(EMB_REPO, name), str)
for name in ("onnx/model_quantized.onnx", "tokenizer.json"))
@unittest.skipUnless(_embedder_cached(), "the embedder is not in the local model cache")
class CachedEmbedderTests(unittest.TestCase):
"""The query embedder resolves its files from the local cache with no
request (rerank.hub_file), and the arena-shrinking RunOptions leave the
vectors bit-identical."""
def test_no_request_and_identical_vectors(self):
from unittest import mock
import numpy as np
from canlex import rerank
from canlex.embed import Embedder
real = rerank.hf_hub_download
calls = []
def spy(*args, **kwargs):
calls.append(kwargs.get("local_files_only", False))
return real(*args, **kwargs)
with mock.patch.object(rerank, "hf_hub_download", side_effect=spy):
emb = Embedder()
self.assertEqual(calls, [True, True])
texts = ["Detention review within 48 hours.", "Seizure of currency."]
with_shrink = emb.encode(texts)
run = emb.session.run
with mock.patch.object(emb.session, "run",
side_effect=lambda names, feed, *a: run(names, feed)):
without = emb.encode(texts)
self.assertTrue(np.array_equal(with_shrink, without))
if __name__ == "__main__":
unittest.main()
|