mirror of
https://github.com/neuml/txtai.git
synced 2026-10-12 02:19:52 +00:00
756 lines
26 KiB
Python
756 lines
26 KiB
Python
"""
|
|
Keyword scoring tests
|
|
"""
|
|
|
|
import os
|
|
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
|
|
from unittest.mock import patch
|
|
|
|
from txtai.scoring import Normalize, ScoringFactory, Scoring
|
|
|
|
|
|
class TestKeyword(unittest.TestCase):
|
|
"""
|
|
Sparse keyword scoring tests.
|
|
"""
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
"""
|
|
Initialize test data.
|
|
"""
|
|
|
|
cls.data = [
|
|
"US tops 5 million confirmed virus cases",
|
|
"Canada's last fully intact ice shelf has suddenly collapsed, forming a Manhattan-sized iceberg",
|
|
"Beijing mobilises invasion craft along coast as Taiwan tensions escalate",
|
|
"The National Park Service warns against sacrificing slower friends in a bear attack",
|
|
"Maine man wins $1M from $25 lottery ticket",
|
|
"wins wins wins",
|
|
"Make huge profits without work, earn up to $100,000 a day",
|
|
]
|
|
|
|
cls.data = [(uid, x, None) for uid, x in enumerate(cls.data)]
|
|
|
|
def testBM25(self):
|
|
"""
|
|
Test bm25
|
|
"""
|
|
|
|
self.runTests("bm25")
|
|
|
|
def testCustom(self):
|
|
"""
|
|
Test custom method
|
|
"""
|
|
|
|
self.runTests("txtai.scoring.BM25")
|
|
|
|
def testCustomInvalid(self):
|
|
"""
|
|
Test invalid custom method
|
|
"""
|
|
|
|
with self.assertRaises(ImportError):
|
|
ScoringFactory.create("pprint.pprint")
|
|
|
|
def testCustomNotFound(self):
|
|
"""
|
|
Test unresolvable custom method
|
|
"""
|
|
|
|
with self.assertRaises(ImportError):
|
|
ScoringFactory.create("notfound.scoring")
|
|
|
|
def testNotImplemented(self):
|
|
"""
|
|
Test exceptions for non-implemented methods
|
|
"""
|
|
|
|
scoring = Scoring()
|
|
|
|
self.assertRaises(NotImplementedError, scoring.insert, None, None)
|
|
self.assertRaises(NotImplementedError, scoring.delete, None)
|
|
self.assertRaises(NotImplementedError, scoring.weights, None)
|
|
self.assertRaises(NotImplementedError, scoring.search, None, None)
|
|
self.assertRaises(NotImplementedError, scoring.batchsearch, None, None, None)
|
|
self.assertRaises(NotImplementedError, scoring.count)
|
|
self.assertRaises(NotImplementedError, scoring.load, None)
|
|
self.assertRaises(NotImplementedError, scoring.save, None)
|
|
self.assertRaises(NotImplementedError, scoring.close)
|
|
self.assertRaises(NotImplementedError, scoring.issparse)
|
|
self.assertRaises(NotImplementedError, scoring.isnormalized)
|
|
self.assertRaises(NotImplementedError, scoring.isbayes)
|
|
|
|
@patch("sqlalchemy.orm.Query.params")
|
|
def testPGText(self, query):
|
|
"""
|
|
Test PGText
|
|
"""
|
|
|
|
# Mock database query
|
|
query.return_value = [(3, 1.0)]
|
|
|
|
# Create scoring
|
|
path = os.path.join(tempfile.gettempdir(), "pgtext.sqlite")
|
|
scoring = ScoringFactory.create({"method": "pgtext", "url": f"sqlite:///{path}", "schema": "txtai"})
|
|
scoring.index((uid, {"text": text}, tags) for uid, text, tags in self.data)
|
|
|
|
# Run search and validate correct result returned
|
|
index, _ = scoring.search("bear", 1)[0]
|
|
self.assertEqual(index, 3)
|
|
|
|
# Run batch search
|
|
index, _ = scoring.batchsearch(["bear"], 1)[0][0]
|
|
self.assertEqual(index, 3)
|
|
|
|
# Validate save/load/delete
|
|
scoring.save(None)
|
|
scoring.load(None)
|
|
|
|
# Validate count
|
|
self.assertEqual(scoring.count(), len(self.data))
|
|
|
|
# Test delete
|
|
scoring.delete([0])
|
|
self.assertEqual(scoring.count(), len(self.data) - 1)
|
|
|
|
# PGText is a normalized sparse index
|
|
self.assertTrue(scoring.issparse() and scoring.isnormalized() and not scoring.isbayes())
|
|
self.assertIsNone(scoring.weights("This is a test".split()))
|
|
|
|
# Close scoring
|
|
scoring.close()
|
|
|
|
def testSIF(self):
|
|
"""
|
|
Test sif
|
|
"""
|
|
|
|
self.runTests("sif")
|
|
|
|
def testTermsEmpty(self):
|
|
"""
|
|
Test searching a terms index with nothing indexed
|
|
"""
|
|
|
|
# No documents ever inserted
|
|
scoring = ScoringFactory.create({"method": "bm25", "terms": True})
|
|
self.assertEqual(scoring.search("bear", 1), [])
|
|
|
|
# Documents present but all skipped (missing text field) - term database never initialized
|
|
scoring = ScoringFactory.create({"method": "bm25", "terms": True})
|
|
scoring.index([(0, {"other": "value"}, None)])
|
|
self.assertEqual(scoring.count(), 0)
|
|
self.assertEqual(scoring.search("bear", 1), [])
|
|
|
|
def testTermsEmptySave(self):
|
|
"""
|
|
Test saving and loading a terms index with nothing indexed
|
|
"""
|
|
|
|
for method in ["bm25", "tfidf"]:
|
|
config = {"method": method, "terms": True}
|
|
|
|
# No documents ever inserted - term database never initialized
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index([])
|
|
|
|
# Save/load and validate index is still empty
|
|
scoring = self.save(scoring, config, f"scoring.{method}.empty")
|
|
self.assertEqual(scoring.count(), 0)
|
|
self.assertEqual(scoring.search("bear", 1), [])
|
|
|
|
def testTermsIdRoundTrip(self):
|
|
"""
|
|
Test that saving and loading preserves keyword document IDs and their content
|
|
"""
|
|
|
|
texts = ["orchard apple", "coastal ocean", "forest trail"]
|
|
cases = [["001", "1", "01"], [0, 1, -1], [1, "1", "notes"], ["²", "³", "¹"], [0, 1, 2], ["one", "two", "three"]]
|
|
for method in ["bm25", "tfidf", "sif"]:
|
|
for content in [False, True]:
|
|
for ids in cases:
|
|
with self.subTest(method=method, content=content, ids=ids), tempfile.TemporaryDirectory() as directory:
|
|
config = {"method": method, "terms": True, "content": content}
|
|
scoring = ScoringFactory.create(config)
|
|
try:
|
|
scoring.index([(uid, text, None) for uid, text in zip(ids, texts)])
|
|
expected = scoring.batchsearch(texts)
|
|
path, copy = os.path.join(directory, "index"), os.path.join(directory, "copy")
|
|
|
|
# Cover an initial save, an in-place save and a copy to a different path.
|
|
for target in [path, path, copy]:
|
|
scoring.save(target)
|
|
scoring.close()
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.load(target)
|
|
self.assertEqual(scoring.batchsearch(texts), expected)
|
|
|
|
# IDs must still identify the same document when deleting after a reload.
|
|
scoring.delete([ids[0]])
|
|
scoring.save(copy)
|
|
scoring.close()
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.load(copy)
|
|
self.assertEqual(scoring.count(), 2)
|
|
self.assertEqual(scoring.search(texts[0]), [])
|
|
self.assertEqual(scoring.batchsearch(texts[1:]), expected[1:])
|
|
finally:
|
|
scoring.close()
|
|
|
|
def testTermsLegacyIds(self):
|
|
"""
|
|
Test reading legacy text IDs and saving additional documents into those indexes
|
|
"""
|
|
|
|
texts = ["orchard apple", "coastal ocean"]
|
|
for ids in [[0, 1], ["one", "two"]]:
|
|
for content in [False, True]:
|
|
with self.subTest(ids=ids, content=content), tempfile.TemporaryDirectory() as directory:
|
|
config = {"method": "bm25", "terms": True, "content": content}
|
|
scoring = ScoringFactory.create(config)
|
|
try:
|
|
scoring.index([(uid, text, None) for uid, text in zip(ids, texts)])
|
|
expected = scoring.batchsearch(texts)
|
|
path = os.path.join(directory, "index")
|
|
scoring.save(path)
|
|
scoring.close()
|
|
|
|
# Older indexes stored IDs as plain SQLite TEXT, without type information.
|
|
with sqlite3.connect(path + ".terms") as connection:
|
|
connection.executemany("UPDATE documents SET id = ? WHERE indexid = ?", [(str(uid), i) for i, uid in enumerate(ids)])
|
|
connection.close()
|
|
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.load(path)
|
|
self.assertEqual(scoring.batchsearch(texts), expected)
|
|
|
|
scoring.upsert([("003", "forest trail", None)])
|
|
expected = scoring.batchsearch([*texts, "forest trail"])
|
|
scoring.save(path)
|
|
scoring.close()
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.load(path)
|
|
self.assertEqual(scoring.count(), 3)
|
|
self.assertEqual(scoring.batchsearch([*texts, "forest trail"]), expected)
|
|
finally:
|
|
scoring.close()
|
|
|
|
def testWildcardEscapedCharacters(self):
|
|
"""
|
|
Test explicit escaping while preserving SQL LIKE wildcard behavior
|
|
"""
|
|
|
|
data = ["report_2026", "reportX2026", "rate%2026", "rateX2026", r"path\_2026", r"path\X2026", "path_2026", "plain2026", "plain2027"]
|
|
queries = [
|
|
("report_*", {0, 1}),
|
|
(r"report\_*", {0}),
|
|
("rate%*", {2, 3}),
|
|
(r"rate\%*", {2}),
|
|
(r"*\%2026", {2}),
|
|
("*_2026", set(range(8))),
|
|
(r"*\_2026", {0, 4, 6}),
|
|
(r"path\\\_*", {4}),
|
|
(r"path\\*", {4, 5}),
|
|
("plain*", {7, 8}),
|
|
("pl**2026", {7}),
|
|
("*2027", {8}),
|
|
("rate%2026", {2}),
|
|
("report_2026", {0}),
|
|
("*", set()),
|
|
]
|
|
for method in ["bm25", "tfidf", "sif"]:
|
|
scoring = ScoringFactory.create({"method": method, "terms": True, "tokenizer": {"whitespace": True}})
|
|
try:
|
|
scoring.index([(uid, text, None) for uid, text in enumerate(data)])
|
|
for query, expected in queries:
|
|
with self.subTest(method=method, query=query):
|
|
results = scoring.search(query, len(data))
|
|
self.assertEqual({uid for uid, _ in results}, expected)
|
|
self.assertEqual(len(results), len(expected))
|
|
finally:
|
|
scoring.close()
|
|
|
|
def testDeleteUnknownId(self):
|
|
"""
|
|
Test that deleting an id that was never indexed is a no-op, not a crash
|
|
"""
|
|
|
|
# Terms index (bm25): Terms.delete used self.ids.index(), raising ValueError
|
|
scoring = ScoringFactory.create({"method": "bm25", "terms": True, "content": True})
|
|
scoring.index(self.data)
|
|
|
|
scoring.delete([0, len(self.data) + 100])
|
|
self.assertFalse(scoring.search("tops", 1))
|
|
|
|
# TFIDF content only, terms disabled: TFIDF.delete used dict.pop(), raising KeyError
|
|
scoring = ScoringFactory.create({"method": "tfidf", "content": True})
|
|
scoring.index(self.data)
|
|
|
|
scoring.delete([0, len(self.data) + 100])
|
|
self.assertNotIn(0, scoring.documents)
|
|
|
|
def testDeleteIdempotent(self):
|
|
"""
|
|
Test that deleting the same id more than once does not corrupt count()
|
|
"""
|
|
|
|
# Terms index (bm25): Terms.delete extended self.deletes without dedup, so a
|
|
# repeated delete appended the same index again and count() undercounted.
|
|
scoring = ScoringFactory.create({"method": "bm25", "terms": True, "content": True})
|
|
scoring.index(self.data)
|
|
|
|
total = scoring.count()
|
|
self.assertEqual(total, len(self.data))
|
|
|
|
scoring.delete([0])
|
|
self.assertEqual(scoring.count(), total - 1)
|
|
|
|
# Deleting the same id again is a no-op for the count
|
|
scoring.delete([0])
|
|
self.assertEqual(scoring.count(), total - 1)
|
|
|
|
def testDeleteMixedTerms(self):
|
|
"""
|
|
Test that common term scoring cannot restore deleted keyword results
|
|
"""
|
|
|
|
for method in ["bm25", "tfidf", "sif"]:
|
|
for content in [False, True]:
|
|
with self.subTest(method=method, content=content), tempfile.TemporaryDirectory() as directory:
|
|
config = {"method": method, "terms": True, "content": content}
|
|
scoring = ScoringFactory.create(config)
|
|
try:
|
|
# rare occurs in one document; common occurs in every document.
|
|
scoring.index([(uid, "rare common" if uid == 0 else "common", None) for uid in range(20)])
|
|
scoring.search("rare common", 20)
|
|
scoring.delete([1])
|
|
|
|
for persisted in [False, True]:
|
|
if persisted:
|
|
path = os.path.join(directory, "scoring")
|
|
scoring.save(path)
|
|
scoring.close()
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.load(path)
|
|
|
|
for query in ["rare", "common", "rare common"]:
|
|
for limit in [3, 20]:
|
|
results = scoring.search(query, limit)
|
|
ids = [result["id"] if content else result[0] for result in results]
|
|
self.assertNotIn(1, ids)
|
|
self.assertEqual(len(ids), 1 if query == "rare" else min(limit, 19))
|
|
if query == "rare common":
|
|
self.assertEqual(ids[0], 0)
|
|
|
|
self.assertEqual(scoring.count(), 19)
|
|
finally:
|
|
scoring.close()
|
|
|
|
def testDeleteReinsert(self):
|
|
"""
|
|
Test that an id deleted, re-added and deleted again is removed from count() and search()
|
|
"""
|
|
|
|
# Terms index: Terms.delete resolved each id with self.ids.index(), which only found the first
|
|
# (already deleted) position. The re-added copy survived, count() was one too high and search()
|
|
# raised a KeyError for the re-added text since the content had been removed.
|
|
for method in ["bm25", "tfidf", "sif"]:
|
|
scoring = ScoringFactory.create({"method": method, "terms": True, "content": True})
|
|
scoring.index(self.data)
|
|
|
|
total = scoring.count()
|
|
|
|
scoring.delete([0])
|
|
self.assertEqual(scoring.count(), total - 1)
|
|
|
|
# Re-add the same id with new text, then delete it again
|
|
scoring.upsert([(0, "Lunar eclipse visible across North America", None)])
|
|
self.assertEqual(scoring.count(), total)
|
|
|
|
scoring.delete([0])
|
|
self.assertFalse(scoring.search("lunar eclipse", 1))
|
|
self.assertEqual(scoring.count(), total - 1)
|
|
|
|
def testTFIDF(self):
|
|
"""
|
|
Test tfidf
|
|
"""
|
|
|
|
self.runTests("tfidf")
|
|
|
|
def runTests(self, method):
|
|
"""
|
|
Runs a series of tests for a scoring method.
|
|
|
|
Args:
|
|
method: scoring method
|
|
"""
|
|
|
|
config = {"method": method}
|
|
|
|
self.index(config)
|
|
self.upsert(config)
|
|
self.weights(config)
|
|
self.search(config)
|
|
self.delete(config)
|
|
self.normalize(config)
|
|
self.content(config)
|
|
self.empty(config)
|
|
self.copy(config)
|
|
self.settings(config)
|
|
self.tokenization(config)
|
|
|
|
def index(self, config, data=None):
|
|
"""
|
|
Test scoring index method.
|
|
|
|
Args:
|
|
config: scoring config
|
|
data: data to index with scoring method
|
|
|
|
Returns:
|
|
scoring
|
|
"""
|
|
|
|
# Derive input data
|
|
data = data if data else self.data
|
|
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index(data)
|
|
|
|
keys = [k for k, v in sorted(scoring.idf.items(), key=lambda x: x[1])]
|
|
|
|
# Test count
|
|
self.assertEqual(scoring.count(), len(data))
|
|
|
|
# Win should be lowest score
|
|
self.assertEqual(keys[0], "wins")
|
|
|
|
# Test save/load
|
|
self.assertIsNotNone(self.save(scoring, config, f"scoring.{config['method']}.index"))
|
|
|
|
# Test search returns none when terms disabled (default)
|
|
self.assertIsNone(scoring.search("query"))
|
|
|
|
return scoring
|
|
|
|
def upsert(self, config):
|
|
"""
|
|
Test scoring upsert method
|
|
"""
|
|
|
|
scoring = ScoringFactory.create({**config, **{"tokenizer": {"alphanum": True, "stopwords": True}}})
|
|
scoring.upsert(self.data)
|
|
|
|
# Test count
|
|
self.assertEqual(scoring.count(), len(self.data))
|
|
|
|
# Test stop word is removed
|
|
self.assertFalse("and" in scoring.idf)
|
|
|
|
def save(self, scoring, config, name):
|
|
"""
|
|
Test scoring index save/load.
|
|
|
|
Args:
|
|
scoring: scoring index
|
|
config: scoring config
|
|
name: output file name
|
|
|
|
Returns:
|
|
scoring
|
|
"""
|
|
|
|
# Generate temp file path
|
|
index = os.path.join(tempfile.gettempdir(), "scoring")
|
|
os.makedirs(index, exist_ok=True)
|
|
|
|
# Save scoring instance
|
|
scoring.save(f"{index}/{name}")
|
|
|
|
# Reload scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.load(f"{index}/{name}")
|
|
|
|
return scoring
|
|
|
|
def weights(self, config):
|
|
"""
|
|
Test standard and tag weighted scores.
|
|
|
|
Args:
|
|
config: scoring config
|
|
"""
|
|
|
|
document = (1, ["bear", "wins"], None)
|
|
|
|
scoring = self.index(config)
|
|
weights = scoring.weights(document[1])
|
|
|
|
# Default weights
|
|
self.assertNotEqual(weights[0], weights[1])
|
|
|
|
data = self.data[:]
|
|
|
|
uid, text, _ = data[3]
|
|
data[3] = (uid, text, "wins")
|
|
|
|
scoring = self.index(config, data)
|
|
weights = scoring.weights(document[1])
|
|
|
|
# Modified weights
|
|
self.assertEqual(weights[0], weights[1])
|
|
|
|
def search(self, config):
|
|
"""
|
|
Test scoring search.
|
|
|
|
Args:
|
|
config: scoring config
|
|
"""
|
|
|
|
# Create combined config
|
|
config = {**config, **{"terms": True}}
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index(self.data)
|
|
|
|
# Run search and validate correct result returned
|
|
index, _ = scoring.search("bear", 1)[0]
|
|
self.assertEqual(index, 3)
|
|
|
|
# Run batch search
|
|
index, _ = scoring.batchsearch(["bear"], 1)[0][0]
|
|
self.assertEqual(index, 3)
|
|
|
|
# Run wildcard search
|
|
index, _ = scoring.search("bea*", 1)[0]
|
|
self.assertEqual(index, 3)
|
|
|
|
# Test save/reload
|
|
self.save(scoring, config, f"scoring.{config['method']}.search")
|
|
|
|
# Re-run search and validate correct result returned
|
|
index, _ = scoring.search("bear", 1)[0]
|
|
self.assertEqual(index, 3)
|
|
|
|
def delete(self, config):
|
|
"""
|
|
Test delete.
|
|
"""
|
|
|
|
# Create combined config
|
|
config = {**config, **{"terms": True, "content": True}}
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index(self.data)
|
|
|
|
# Run search and validate correct result returned
|
|
index = scoring.search("bear", 1)[0]["id"]
|
|
|
|
# Delete result and validate the query no longer returns results
|
|
scoring.delete([index])
|
|
self.assertFalse(scoring.search("bear", 1))
|
|
|
|
# Save and validate count
|
|
self.save(scoring, config, f"scoring.{config['method']}.delete")
|
|
self.assertEqual(scoring.count(), len(self.data) - 1)
|
|
|
|
def normalize(self, config):
|
|
"""
|
|
Test scoring search with normalized scores.
|
|
|
|
Args:
|
|
method: scoring method
|
|
"""
|
|
|
|
# Default normalization
|
|
scoring = ScoringFactory.create({**config, **{"terms": True, "normalize": True}})
|
|
scoring.index(self.data)
|
|
|
|
# Run search and validate correct result returned
|
|
index, score = scoring.search(self.data[3][1], 1)[0]
|
|
self.assertEqual(index, 3)
|
|
self.assertEqual(score, 1.0)
|
|
|
|
# Bayesian normalization with default dynamic alpha/beta settings
|
|
baseline = ScoringFactory.create({**config, **{"terms": True}})
|
|
baseline.index(self.data)
|
|
|
|
scoring = ScoringFactory.create({**config, **{"terms": True, "normalize": "bayes"}})
|
|
scoring.index(self.data)
|
|
|
|
query = "wins"
|
|
base = baseline.search(query, 3)
|
|
bayes = scoring.search(query, 3)
|
|
|
|
# Bayesian normalization should preserve ranking order while mapping scores to [0, 1]
|
|
self.assertEqual([uid for uid, _ in base], [uid for uid, _ in bayes])
|
|
self.assertTrue(all(0.0 <= score <= 1.0 for _, score in bayes))
|
|
|
|
# BB25 alias should resolve to Bayesian normalization
|
|
scoring = ScoringFactory.create({**config, **{"terms": True, "normalize": "bb25"}})
|
|
scoring.index(self.data)
|
|
bb25 = scoring.search(query, 3)
|
|
self.assertEqual([uid for uid, _ in base], [uid for uid, _ in bb25])
|
|
self.assertTrue(all(0.0 <= score <= 1.0 for _, score in bb25))
|
|
|
|
# BB25 candidate-set behavior: zero scores remain 0, positive scores are transformed
|
|
normalizer = Normalize("bb25")
|
|
scores = normalizer([(0, 0.0), (1, 1.0), (2, 2.0)], scoring.avgscore)
|
|
self.assertEqual(scores[0][1], 0.0)
|
|
self.assertGreater(scores[1][1], 0.0)
|
|
self.assertGreater(scores[2][1], scores[1][1])
|
|
|
|
# Test negative scores
|
|
scores = normalizer([(0, -100.0)], scoring.avgscore)
|
|
self.assertEqual(scores[0][1], 0.0)
|
|
|
|
# Bayesian normalization with custom parameters
|
|
config = {**config, **{"terms": True, "normalize": {"method": "bayes", "alpha": 2.0}}}
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index(self.data)
|
|
|
|
custom = scoring.search(query, 3)
|
|
self.assertEqual([uid for uid, _ in base], [uid for uid, _ in custom])
|
|
self.assertTrue(all(0.0 <= score <= 1.0 for _, score in custom))
|
|
|
|
def content(self, config):
|
|
"""
|
|
Test scoring search with content.
|
|
|
|
Args:
|
|
config: scoring config
|
|
"""
|
|
|
|
scoring = ScoringFactory.create({**config, **{"terms": True, "content": True}})
|
|
scoring.index(self.data)
|
|
|
|
# Test text with content
|
|
text = "Great news today"
|
|
scoring.index([(scoring.total, text, None)])
|
|
|
|
# Run search and validate correct result returned
|
|
result = scoring.search("great news", 1)[0]["text"]
|
|
self.assertEqual(result, text)
|
|
|
|
# Test reading text from dictionary
|
|
text = "Feel good story: baby panda born"
|
|
scoring.index([(scoring.total, {"text": text}, None)])
|
|
|
|
# Run search and validate correct result returned
|
|
result = scoring.search("feel good story", 1)[0]["text"]
|
|
self.assertEqual(result, text)
|
|
|
|
def empty(self, config):
|
|
"""
|
|
Test scoring index properly handles an index call when no data present.
|
|
|
|
Args:
|
|
config: scoring config
|
|
"""
|
|
|
|
# Create scoring index with no data
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index([])
|
|
|
|
# Assert index call returns and index has a count of 0
|
|
self.assertEqual(scoring.total, 0)
|
|
|
|
def copy(self, config):
|
|
"""
|
|
Test scoring index copy method.
|
|
"""
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create({**config, **{"terms": True}})
|
|
scoring.index(self.data)
|
|
|
|
# Generate temp file path
|
|
index = os.path.join(tempfile.gettempdir(), "scoring")
|
|
os.makedirs(index, exist_ok=True)
|
|
|
|
# Create file to test replacing existing file
|
|
path = f"{index}/scoring.{config['method']}.copy"
|
|
with open(f"{index}.terms", "w", encoding="utf-8") as f:
|
|
f.write("TEST")
|
|
|
|
# Save scoring instance
|
|
scoring.save(path)
|
|
self.assertTrue(os.path.exists(path))
|
|
|
|
@patch("sys.byteorder", "big")
|
|
def settings(self, config):
|
|
"""
|
|
Test various settings.
|
|
|
|
Args:
|
|
config: scoring config
|
|
"""
|
|
|
|
# Create combined config
|
|
config = {**config, **{"terms": {"cachelimit": 0, "cutoff": 0.25, "wal": True}}}
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index(self.data)
|
|
|
|
# Save/load index
|
|
self.save(scoring, config, f"scoring.{config['method']}.settings")
|
|
|
|
index, _ = scoring.search("bear bear bear wins", 1)[0]
|
|
self.assertEqual(index, 3)
|
|
|
|
# Save to same path
|
|
self.save(scoring, config, f"scoring.{config['method']}.settings")
|
|
|
|
# Save to different path
|
|
self.save(scoring, config, f"scoring.{config['method']}.move")
|
|
|
|
# Validate counts
|
|
self.assertEqual(scoring.count(), len(self.data))
|
|
|
|
def tokenization(self, config):
|
|
"""
|
|
Test tokenization methods.
|
|
|
|
Args:
|
|
config: scoring config
|
|
"""
|
|
|
|
# Test whitespace tokenization
|
|
config = {**config, **{"terms": True, "tokenizer": {"whitespace": True}}}
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index([(0, "abc-def-123", None)])
|
|
|
|
self.assertEqual(scoring.search("abc-def-123")[0][0], 0)
|
|
|
|
# Test regular expression tokenization
|
|
config = {**config, **{"tokenizer": {"regexp": r"\w{5,}"}}}
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index([(0, "hello test", None)])
|
|
|
|
self.assertEqual(scoring.search("hello")[0][0], 0)
|
|
self.assertFalse(scoring.search("test"))
|
|
|
|
# Test ngram tokenization
|
|
ngrams = {"ngrams": 3, "lpad": " ", "rpad": " ", "unique": True}
|
|
config = {**config, **{"tokenizer": {"ngrams": ngrams}}}
|
|
|
|
# Create scoring instance
|
|
scoring = ScoringFactory.create(config)
|
|
scoring.index([(0, "hello test", None)])
|
|
|
|
self.assertEqual(scoring.search("hello")[0][0], 0)
|