Spaces:
Sleeping
Sleeping
| """sentence-transformers backend (the default, BGE-small on a normal machine), exercised with a tiny | |
| locally-built BERT so the test needs no download.""" | |
| import re | |
| import numpy as np | |
| import pytest | |
| from mrag.config import Config | |
| def tiny_st(tmp_path_factory): | |
| pytest.importorskip("sentence_transformers") | |
| from sentence_transformers import SentenceTransformer, models | |
| from transformers import BertConfig, BertModel, BertTokenizerFast | |
| d = tmp_path_factory.mktemp("tiny") | |
| words = sorted(set(re.findall(r"[a-z]+", "the warranty lasts thirty six months for every robot unit"))) | |
| (d / "vocab.txt").write_text("\n".join(["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + words)) | |
| tok = BertTokenizerFast(str(d / "vocab.txt")) | |
| m = BertModel(BertConfig(vocab_size=len(words) + 5, hidden_size=32, num_hidden_layers=1, | |
| num_attention_heads=2, intermediate_size=64)) | |
| tok.save_pretrained(d / "hf") | |
| m.save_pretrained(d / "hf") | |
| st = SentenceTransformer(modules=[models.Transformer(str(d / "hf")), models.Pooling(32, "cls"), models.Normalize()]) | |
| st.save(str(d / "st")) | |
| return str(d / "st") | |
| def test_sentence_transformers_backend(tiny_st): | |
| from mrag.models import TextEmbedder | |
| e = TextEmbedder(Config(text_backend="sentence-transformers", text_model=tiny_st, query_prefix="query: ")) | |
| D = e.embed_docs(["the warranty lasts", "every robot unit"]) | |
| q = e.embed_query("warranty") | |
| assert e.dim == 32 and D.shape == (2, 32) and q.shape == (32,) | |
| assert np.allclose(np.linalg.norm(D, axis=1), 1, atol=1e-5) and abs(np.linalg.norm(q) - 1) < 1e-5 | |
| assert e.name == f"st:{tiny_st}" | |
| def test_wordllama_backend(): | |
| from mrag.models import TextEmbedder | |
| e = TextEmbedder(Config(text_backend="wordllama")) | |
| v = e.embed_docs(["The cat sat on the mat.", "A kitten is resting on a rug.", "Quarterly revenue grew."]) | |
| sims = v @ v.T | |
| assert sims[0, 1] > sims[0, 2] + 0.2 # semantic, not lexical: cat~kitten share no words | |
| def test_unknown_backend(): | |
| from mrag.models import TextEmbedder | |
| with pytest.raises(ValueError): | |
| TextEmbedder(Config(text_backend="nope")) | |