def test_multiple_fields(self): from corpus import InMemoryDocument, InMemoryCorpus from suffixarray import SuffixArray corpus = InMemoryCorpus() corpus.add_document( InMemoryDocument(0, { "field1": "a b c", "field2": "b c d" })) corpus.add_document(InMemoryDocument(1, { "field1": "x", "field2": "y" })) corpus.add_document(InMemoryDocument(2, { "field1": "y", "field2": "z" })) engine0 = SuffixArray(corpus, ["field1", "field2"], self._normalizer, self._tokenizer) engine1 = SuffixArray(corpus, ["field1"], self._normalizer, self._tokenizer) engine2 = SuffixArray(corpus, ["field2"], self._normalizer, self._tokenizer) self._process_query_and_verify_winner(engine0, "b c", [0], 2) self._process_query_and_verify_winner(engine0, "y", [1, 2], 1) self._process_query_and_verify_winner(engine1, "x", [1], 1) self._process_query_and_verify_winner(engine1, "y", [2], 1) self._process_query_and_verify_winner(engine1, "z", [], None) self._process_query_and_verify_winner(engine2, "z", [2], 1)
def populateSuffixArray(keys): """ Takes a list of keys and initializes and populates a suffixarray.SuffixArray object """ suffix_array = SuffixArray() for key in keys: suffix_array.insert(key) return suffix_array
def populateSuffixArray(keys): """ Takes a list of keys and initializes and populates a suffixarray.SuffixArray object """ suffix_array = SuffixArray() for key in keys: suffix_array.insert(key) return suffix_array
def test_memory_usage(self): import tracemalloc import inspect from corpus import InMemoryDocument, InMemoryCorpus from suffixarray import SuffixArray corpus = InMemoryCorpus() corpus.add_document( InMemoryDocument(0, { "a": "o o\n\n\no\n\no", "b": "o o\no \no" })) corpus.add_document(InMemoryDocument(1, {"a": "ba", "b": "b bab"})) corpus.add_document(InMemoryDocument(2, {"a": "o o O o", "b": "o o"})) corpus.add_document(InMemoryDocument(3, {"a": "oO" * 10000, "b": "o"})) corpus.add_document( InMemoryDocument(4, { "a": "cbab o obab O ", "b": "o o " * 10000 })) tracemalloc.start() snapshot1 = tracemalloc.take_snapshot() engine = SuffixArray(corpus, ["a", "b"], self._normalizer, self._tokenizer) snapshot2 = tracemalloc.take_snapshot() tracemalloc.stop() for statistic in snapshot2.compare_to(snapshot1, "filename"): if statistic.traceback[0].filename == inspect.getfile(SuffixArray): self.assertLessEqual(statistic.size_diff, 2000000, "Memory usage seems excessive.")
def test_cran_corpus(self): import os.path from corpus import InMemoryCorpus from suffixarray import SuffixArray corpus = InMemoryCorpus(os.path.join(data_path, 'cran.xml')) engine = SuffixArray(corpus, ["body"], self._normalizer, self._tokenizer) self._process_query_and_verify_winner(engine, "visc", [328], 11) self._process_query_and_verify_winner(engine, "Of A", [946], 10) self._process_query_and_verify_winner(engine, "", [], None) self._process_query_and_verify_winner(engine, "approximate solution", [159, 1374], 3)
def assignment_b_suffixarray_1(): # Use these throughout below. normalizer = BrainDeadNormalizer() tokenizer = BrainDeadTokenizer() # Prepare for some suffix array lookups. print("LOADING...") corpus = InMemoryCorpus("data/cran.xml") print("INDEXING...") engine = SuffixArray(corpus, ["body"], normalizer, tokenizer) results = [] hit_count = 5 # Callback for receiving matches. def match_collector(match): results.append(match) print("*** WINNER", match["score"], match["document"]) # Define the actual test queries. test1 = ("visc", 11, [328]) # Look for {'viscous', 'viscosity', ...}. test2 = ("Of A", 10, [946]) # Test robustness for case and whitespace. test3 = ("", 0, []) # Safety feature: Match nothing instead of everything. test4 = ("approximate solution", 3, [1374, 159]) # Multiple winners. # Test that the simple occurrence ranking works. Be robust towards how ties are resolved. for (query, winner_score, winner_document_ids) in [test1, test2, test3, test4]: print("SEARCHING for '" + query + "'...") results.clear() engine.evaluate(query, { "debug": False, "hit_count": hit_count }, match_collector) if winner_document_ids: assert results[0]["score"] == winner_score assert results[0]["document"].document_id in winner_document_ids assert len(results) <= hit_count else: assert len(results) == 0
def main(): import os.path from normalization import BrainDeadNormalizer from tokenization import BrainDeadTokenizer from corpus import InMemoryCorpus from suffixarray import SuffixArray print("Building suffix array from Cranfield corpus...") normalizer = BrainDeadNormalizer() tokenizer = BrainDeadTokenizer() corpus = InMemoryCorpus(os.path.join(data_path, 'cran.xml')) engine = SuffixArray(corpus, ["body"], normalizer, tokenizer) options = {"debug": False, "hit_count": 5} print("Enter a prefix phrase query and find matching documents.") print(f"Lookup options are {options}.") print("Returned scores are occurrence counts.") def evaluator(query): matches = [] engine.evaluate(query, options, lambda m: matches.append(m)) return matches simple_repl("query", evaluator)
def assignment_b_suffixarray_2(): # For testing. class TestNormalizer(Normalizer): _table = str.maketrans({'Ø': 'O'}) def canonicalize(self, buffer: str) -> str: return buffer def normalize(self, token: str) -> str: return token.upper().translate(self._table) # For testing. class TestDocument(Document): def __init__(self, document_id: int, a: str, b: str): self._document_id = document_id self._a = a self._b = b def get_document_id(self) -> int: return self._document_id def get_field(self, field_name: str, default: str) -> str: if field_name == "a": return self._a if field_name == "b": return self._b return default # For testing. class TestCorpus(Corpus): def __init__(self): self._docs = [] self._docs.append(TestDocument(len(self._docs), "ø o\n\n\nø\n\no", "ø o\nø \no")) self._docs.append(TestDocument(len(self._docs), "ba", "b bab")) self._docs.append(TestDocument(len(self._docs), "ø o Ø o", "ø o")) self._docs.append(TestDocument(len(self._docs), "øO" * 10000, "o")) self._docs.append(TestDocument(len(self._docs), "cbab o øbab Ø ", "ø o " * 10000)) def __iter__(self): return iter(self._docs) def size(self) -> int: return len(self._docs) def get_document(self, document_id: int) -> Document: return self._docs[document_id] # Run the tests! for fields in [("b",), ("a", "b")]: # Create the suffix array over the given set of fields. Measure memory usage. If memory usage is # excessive, most likely the implementation is copying strings or doing other silly stuff instead # of working with buffer indices. The naive reference implementation is not in any way optimized, # and uses about 1.5 MB of memory on this corpus. tracemalloc.start() snapshot1 = tracemalloc.take_snapshot() engine = SuffixArray(TestCorpus(), fields, TestNormalizer(), BrainDeadTokenizer()) snapshot2 = tracemalloc.take_snapshot() for statistic in snapshot2.compare_to(snapshot1, "filename"): if statistic.traceback[0].filename == inspect.getfile(SuffixArray): assert statistic.size_diff < 2000000, f"Memory usage is {statistic.size_diff}" tracemalloc.stop() results = [] def process(m): results.append((m['document'].document_id, m['score'])) expected_results = { ('b',): ( ('bab', [(1, 1)]), ('ø o', [(4, 19999), (0, 3), (2, 1)]), ('o O', [(4, 19999), (0, 3), (2, 1)]), ('oooooo', []), ('o o o o', [(4, 19997), (0, 1)]), ), ('a', 'b'): ( ('bab', [(1, 1)]), ('ø o', [(4, 20000), (0, 6), (2, 4)]), ('o O', [(4, 20000), (0, 6), (2, 4)]), ('oøØOøO', [(3, 1), ]), ('o o o o', [(4, 19997), (0, 2), (2, 1)]), ) } for query, expected in expected_results[fields]: results.clear() engine.evaluate(query, {'hit_count': 10}, process) assert results == expected
def assignment_b_suffixarray_2(): # For testing. class TestNormalizer(Normalizer): _table = str.maketrans({'Ø': 'O'}) def canonicalize(self, buffer: str) -> str: return buffer def normalize(self, token: str) -> str: return token.upper().translate(self._table) # For testing. class TestDocument(Document): def __init__(self, document_id: int, a: str, b: str): self._document_id = document_id self._a = a self._b = b def get_document_id(self) -> int: return self._document_id def get_field(self, field_name: str, default: str) -> str: if field_name == "a": return self._a if field_name == "b": return self._b return default # For testing. class TestCorpus(Corpus): def __init__(self): self._docs = [] self._docs.append( TestDocument(len(self._docs), "ø o\n\n\nø\n\no", "ø o\nø \no")) self._docs.append(TestDocument(len(self._docs), "ba", "b bab")) self._docs.append(TestDocument(len(self._docs), "ø o Ø o", "ø o")) self._docs.append(TestDocument(len(self._docs), "øO" * 10000, "o")) self._docs.append( TestDocument(len(self._docs), "cbab o øbab Ø ", "ø o " * 10000)) def __iter__(self): return iter(self._docs) def size(self) -> int: return len(self._docs) def get_document(self, document_id: int) -> Document: return self._docs[document_id] # Run the tests! for fields in [("b", ), ("a", "b")]: engine = SuffixArray(TestCorpus(), fields, TestNormalizer(), BrainDeadTokenizer()) results = [] def process(m): results.append((m['document'].document_id, m['score'])) expected_results = { ('b', ): ( ('bab', [(1, 1)]), ('ø o', [(4, 19999), (0, 3), (2, 1)]), ('o O', [(4, 19999), (0, 3), (2, 1)]), ('oooooo', []), ('o o o o', [(4, 19997), (0, 1)]), ), ('a', 'b'): ( ('bab', [(1, 1)]), ('ø o', [(4, 20000), (0, 6), (2, 4)]), ('o O', [(4, 20000), (0, 6), (2, 4)]), ('oøØOøO', [ (3, 1), ]), ('o o o o', [(4, 19997), (0, 2), (2, 1)]), ) } for query, expected in expected_results[fields]: results.clear() engine.evaluate(query, {'hit_count': 10}, process) assert results == expected
def assignment_b(): # Use these throughout below. normalizer = BrainDeadNormalizer() tokenizer = BrainDeadTokenizer() # Prepare for some suffix array lookups. print("LOADING...") corpus = InMemoryCorpus("data/cran.xml") print("INDEXING...") engine = SuffixArray(corpus, ["body"], normalizer, tokenizer) results = [] hit_count = 5 # Callback for receiving matches. def match_collector(match): results.append(match) print("*** WINNER", match["score"], match["document"]) # Define the actual test queries. test1 = ("visc", 11, [328]) # Look for {'viscous', 'viscosity', ...}. test2 = ("Of A", 10, [946]) # Test robustness for case and whitespace. test3 = ("", 0, []) # Safety feature: Match nothing instead of everything. test4 = ("approximate solution", 3, [1374, 159]) # Multiple winners. # Test that the simple occurrence ranking works. Be robust towards how ties are resolved. for (query, winner_score, winner_document_ids) in [test1, test2, test3, test4]: print("SEARCHING for '" + query + "'...") results.clear() engine.evaluate(query, { "debug": False, "hit_count": hit_count }, match_collector) if winner_document_ids: assert results[0]["score"] == winner_score assert results[0]["document"].document_id in winner_document_ids assert len(results) <= hit_count else: assert len(results) == 0 # Simple test of using a trie-encoded dictionary for efficiently locating substrings in a buffer. trie = Trie() for s in [ "romerike", "apple computer", "norsk", "norsk ørret", "sverige", "ørret", "banan" ]: trie.add(s, tokenizer) finder = StringFinder(trie, tokenizer) buffer = "det var en gang en norsk ørret fra romerike som likte abba fra sverige" print("SCANNING...") results.clear() finder.scan(buffer, lambda m: results.append(m)) print("Buffer \"" + buffer + "\" contains", results) assert [m["match"] for m in results ] == ["norsk", "norsk ørret", "ørret", "romerike", "sverige"] # Find all MeSH terms that occur verbatim in some selected Cranfield documents! Since MeSH # documents are medical terms and the Cranfield documents have technical content, the # overlap probably isn't that big. print("LOADING...") mesh = InMemoryCorpus("data/mesh.txt") cranfield = InMemoryCorpus("data/cran.xml") print("BUILDING...") trie = Trie() for d in mesh: trie.add(d["body"] or "", tokenizer) finder = StringFinder(trie, tokenizer) print("SCANNING...") for (document_id, expected_matches) in [(0, ["wing", "wing"]), (3, ["solutions", "skin", "friction"]), (1254, ["electrons", "ions"])]: document = cranfield[document_id] buffer = document["body"] or "" results.clear() finder.scan(buffer, lambda m: results.append(m)) print("Cranfield document", document, "contains MeSH terms", results) assert [m["match"] for m in results] == expected_matches