def test_ngrams():
"""Test ngram limits in mention extraction"""
file_name = "lincoln_short"
docs_path = f"tests/data/pure_html/{file_name}.html"
doc = parse_doc(docs_path, file_name)
# Mention Extraction
Person = mention_subclass("Person")
person_ngrams = MentionNgrams(n_max=3)
person_matcher = PersonMatcher()
mention_extractor_udf = MentionExtractorUDF(
[Person], [person_ngrams], [person_matcher]
)
doc = mention_extractor_udf.apply(doc)
assert len(doc.persons) == 118
mentions = doc.persons
assert len([x for x in mentions if x.context.get_num_words() == 1]) == 49
assert len([x for x in mentions if x.context.get_num_words() > 3]) == 0
# Test for unigram exclusion
for mention in doc.persons[:]:
doc.persons.remove(mention)
assert len(doc.persons) == 0
person_ngrams = MentionNgrams(n_min=2, n_max=3)
mention_extractor_udf = MentionExtractorUDF(
[Person], [person_ngrams], [person_matcher]
)
doc = mention_extractor_udf.apply(doc)
assert len(doc.persons) == 69
mentions = doc.persons
assert len([x for x in mentions if x.context.get_num_words() == 1]) == 0
assert len([x for x in mentions if x.context.get_num_words() > 3]) == 0
def test_mention_longest_match():
"""Test longest match filtering in mention extraction."""
file_name = "lincoln_short"
docs_path = f"tests/data/pure_html/{file_name}.html"
doc = parse_doc(docs_path, file_name)
# Mention Extraction
name_ngrams = MentionNgramsPart(n_max=3)
place_ngrams = MentionNgramsTemp(n_max=4)
Name = mention_subclass("Name")
Place = mention_subclass("Place")
def is_birthplace_table_row(mention):
if not mention.sentence.is_tabular():
return False
ngrams = get_row_ngrams(mention, lower=True)
if "birth_place" in ngrams:
return True
else:
return False
birthplace_matcher = LambdaFunctionMatcher(
func=is_birthplace_table_row, longest_match_only=False
)
mention_extractor_udf = MentionExtractorUDF(
[Name, Place],
[name_ngrams, place_ngrams],
[PersonMatcher(), birthplace_matcher],
)
doc = mention_extractor_udf.apply(doc)
mentions = doc.places
mention_spans = [x.context.get_span() for x in mentions]
assert "Sinking Spring Farm" in mention_spans
assert "Farm" in mention_spans
assert len(mention_spans) == 23
# Clear manually
for mention in doc.places[:]:
doc.places.remove(mention)
birthplace_matcher = LambdaFunctionMatcher(
func=is_birthplace_table_row, longest_match_only=True
)
mention_extractor_udf = MentionExtractorUDF(
[Name, Place],
[name_ngrams, place_ngrams],
[PersonMatcher(), birthplace_matcher],
)
doc = mention_extractor_udf.apply(doc)
mentions = doc.places
mention_spans = [x.context.get_span() for x in mentions]
assert "Sinking Spring Farm" in mention_spans
assert "Farm" not in mention_spans
assert len(mention_spans) == 4
def test_row_col_ngram_extraction():
"""Test whether row/column ngrams list is empty, if mention is not in a table."""
file_name = "lincoln_short"
docs_path = f"tests/data/pure_html/{file_name}.html"
doc = parse_doc(docs_path, file_name)
# Mention Extraction
place_ngrams = MentionNgramsTemp(n_max=4)
Place = mention_subclass("Place")
def get_row_and_column_ngrams(mention):
row_ngrams = list(get_row_ngrams(mention))
col_ngrams = list(get_col_ngrams(mention))
if not mention.sentence.is_tabular():
assert len(row_ngrams) == 1 and row_ngrams[0] is None
assert len(col_ngrams) == 1 and col_ngrams[0] is None
else:
assert not any(x is None for x in row_ngrams)
assert not any(x is None for x in col_ngrams)
if "birth_place" in row_ngrams:
return True
else:
return False
birthplace_matcher = LambdaFunctionMatcher(func=get_row_and_column_ngrams)
mention_extractor_udf = MentionExtractorUDF(
[Place], [place_ngrams], [birthplace_matcher]
)
doc = mention_extractor_udf.apply(doc)
def test_candidate_with_nullable_mentions():
"""Test if mentions can be NULL."""
docs_path = "tests/data/html/112823.html"
pdf_path = "tests/data/pdf/"
doc = parse_doc(docs_path, "112823", pdf_path)
# Mention Extraction
MentionTemp = mention_subclass("MentionTemp")
temp_ngrams = MentionNgramsTemp(n_max=2)
mention_extractor_udf = MentionExtractorUDF(
[MentionTemp],
[temp_ngrams],
[temp_matcher],
)
doc = mention_extractor_udf.apply(doc)
assert len(doc.mention_temps) == 23
# Candidate Extraction
CandidateTemp = candidate_subclass("CandidateTemp", [MentionTemp],
nullables=[True])
candidate_extractor_udf = CandidateExtractorUDF([CandidateTemp], [None],
False, False, True)
doc = candidate_extractor_udf.apply(doc, split=0)
# The number of extracted candidates should be that of mentions + 1 (NULL)
assert len(doc.candidate_temps) == len(doc.mention_temps) + 1
# Extracted candidates should include one with NULL mention.
assert None in [c[0] for c in doc.candidate_temps]
def test_visualizer():
from fonduer.utils.visualizer import Visualizer # noqa
"""Unit test of visualizer using the md document.
"""
docs_path = "tests/data/html_simple/md.html"
pdf_path = "tests/data/pdf_simple/md.pdf"
# Grab the md document
doc = parse_doc(docs_path, "md", pdf_path)
assert doc.name == "md"
organization_ngrams = MentionNgrams(n_max=1)
Org = mention_subclass("Org")
organization_matcher = OrganizationMatcher()
mention_extractor_udf = MentionExtractorUDF([Org], [organization_ngrams],
[organization_matcher])
doc = mention_extractor_udf.apply(doc)
Organization = candidate_subclass("Organization", [Org])
candidate_extractor_udf = CandidateExtractorUDF([Organization], None,
False, False, True)
doc = candidate_extractor_udf.apply(doc, split=0)
cands = doc.organizations
# Test visualizer
pdf_path = "tests/data/pdf_simple"
vis = Visualizer(pdf_path)
vis.display_candidates([cands[0]])
def test_visualizer():
"""Unit test of visualizer using the md document."""
from fonduer.utils.visualizer import Visualizer, get_box # noqa
docs_path = "tests/data/html_simple/md.html"
pdf_path = "tests/data/pdf_simple/"
# Grab the md document
doc = parse_doc(docs_path, "md", pdf_path)
assert doc.name == "md"
organization_ngrams = MentionNgrams(n_max=1)
Org = mention_subclass("Org")
organization_matcher = OrganizationMatcher()
mention_extractor_udf = MentionExtractorUDF([Org], [organization_ngrams],
[organization_matcher])
doc = mention_extractor_udf.apply(doc)
Organization = candidate_subclass("Organization", [Org])
candidate_extractor_udf = CandidateExtractorUDF([Organization], None,
False, False, True)
doc = candidate_extractor_udf.apply(doc, split=0)
# Take one candidate
cand = doc.organizations[0]
pdf_path = "tests/data/pdf_simple"
vis = Visualizer(pdf_path)
# Test bounding boxes
boxes = [get_box(mention.context) for mention in cand.get_mentions()]
for box in boxes:
assert box.top <= box.bottom
assert box.left <= box.right
assert boxes == [
mention.context.get_bbox() for mention in cand.get_mentions()
]
# Test visualizer
vis.display_candidates([cand])
def test_multimodal_cand():
"""Test multimodal candidate generation"""
file_name = "radiology"
docs_path = f"tests/data/pure_html/{file_name}.html"
doc = parse_doc(docs_path, file_name)
assert len(doc.sentences) == 35
# Mention Extraction
ms_doc = mention_subclass("m_doc")
ms_sec = mention_subclass("m_sec")
ms_tab = mention_subclass("m_tab")
ms_fig = mention_subclass("m_fig")
ms_cell = mention_subclass("m_cell")
ms_para = mention_subclass("m_para")
ms_cap = mention_subclass("m_cap")
ms_sent = mention_subclass("m_sent")
m_doc = MentionDocuments()
m_sec = MentionSections()
m_tab = MentionTables()
m_fig = MentionFigures()
m_cell = MentionCells()
m_para = MentionParagraphs()
m_cap = MentionCaptions()
m_sent = MentionSentences()
ms = [ms_doc, ms_cap, ms_sec, ms_tab, ms_fig, ms_para, ms_sent, ms_cell]
m = [m_doc, m_cap, m_sec, m_tab, m_fig, m_para, m_sent, m_cell]
matchers = [DoNothingMatcher()] * 8
mention_extractor_udf = MentionExtractorUDF(ms, m, matchers)
doc = mention_extractor_udf.apply(doc)
assert len(doc.m_docs) == 1
assert len(doc.m_caps) == 2
assert len(doc.m_secs) == 5
assert len(doc.m_tabs) == 2
assert len(doc.m_figs) == 2
assert len(doc.m_paras) == 30
assert len(doc.m_sents) == 35
assert len(doc.m_cells) == 21
def test_multinary_relation_feature_extraction():
"""Test extracting candidates from mentions from documents."""
docs_path = "tests/data/html/112823.html"
pdf_path = "tests/data/pdf/112823.pdf"
# Parsing
doc = parse_doc(docs_path, "112823", pdf_path)
assert len(doc.sentences) == 799
# Mention Extraction
part_ngrams = MentionNgrams(n_max=1)
temp_ngrams = MentionNgrams(n_max=1)
volt_ngrams = MentionNgrams(n_max=1)
Part = mention_subclass("Part")
Temp = mention_subclass("Temp")
Volt = mention_subclass("Volt")
mention_extractor_udf = MentionExtractorUDF(
[Part, Temp, Volt],
[part_ngrams, temp_ngrams, volt_ngrams],
[part_matcher, temp_matcher, volt_matcher],
)
doc = mention_extractor_udf.apply(doc)
assert len(doc.parts) == 62
assert len(doc.temps) == 16
assert len(doc.volts) == 33
part = doc.parts[0]
temp = doc.temps[0]
volt = doc.volts[0]
logger.info(f"Part: {part.context}")
logger.info(f"Temp: {temp.context}")
logger.info(f"Volt: {volt.context}")
# Candidate Extraction
PartTempVolt = candidate_subclass("PartTempVolt", [Part, Temp, Volt])
candidate_extractor_udf = CandidateExtractorUDF([PartTempVolt], None,
False, False, True)
doc = candidate_extractor_udf.apply(doc, split=0)
# Manually set id as it is not set automatically b/c a database is not used.
i = 0
for cand in doc.part_temp_volts:
cand.id = i
i = i + 1
n_cands = len(doc.part_temp_volts)
# Featurization based on default feature library
featurizer_udf = FeaturizerUDF([PartTempVolt], FeatureExtractor())
# Test that featurization default feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_default_feats = len(key_set)
# Example feature extractor
def feat_ext(candidates):
candidates = candidates if isinstance(candidates,
list) else [candidates]
for candidate in candidates:
yield candidate.id, f"cand_id_{candidate.id}", 1
# Featurization with one extra feature extractor
feature_extractors = FeatureExtractor(customize_feature_funcs=[feat_ext])
featurizer_udf = FeaturizerUDF([PartTempVolt],
feature_extractors=feature_extractors)
# Test that featurization default feature library with one extra feature extractor
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_default_w_customized_features = len(key_set)
# Example spurious feature extractor
def bad_feat_ext(candidates):
raise RuntimeError()
# Featurization with a spurious feature extractor
feature_extractors = FeatureExtractor(
customize_feature_funcs=[bad_feat_ext])
featurizer_udf = FeaturizerUDF([PartTempVolt],
feature_extractors=feature_extractors)
# Test that featurization default feature library with one extra feature extractor
logger.info("Featurizing with a spurious feature extractor...")
with pytest.raises(RuntimeError):
features = featurizer_udf.apply(doc)
# Featurization with only textual feature
feature_extractors = FeatureExtractor(features=["textual"])
featurizer_udf = FeaturizerUDF([PartTempVolt],
feature_extractors=feature_extractors)
# Test that featurization textual feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_textual_features = len(key_set)
# Featurization with only tabular feature
feature_extractors = FeatureExtractor(features=["tabular"])
featurizer_udf = FeaturizerUDF([PartTempVolt],
feature_extractors=feature_extractors)
# Test that featurization tabular feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_tabular_features = len(key_set)
# Featurization with only structural feature
feature_extractors = FeatureExtractor(features=["structural"])
featurizer_udf = FeaturizerUDF([PartTempVolt],
feature_extractors=feature_extractors)
# Test that featurization structural feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_structural_features = len(key_set)
# Featurization with only visual feature
feature_extractors = FeatureExtractor(features=["visual"])
featurizer_udf = FeaturizerUDF([PartTempVolt],
feature_extractors=feature_extractors)
# Test that featurization visual feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_visual_features = len(key_set)
assert (n_default_feats == n_textual_features + n_tabular_features +
n_structural_features + n_visual_features)
assert n_default_w_customized_features == n_default_feats + n_cands
def test_unary_relation_feature_extraction():
"""Test extracting unary candidates from mentions from documents."""
docs_path = "tests/data/html/112823.html"
pdf_path = "tests/data/pdf/112823.pdf"
# Parsing
doc = parse_doc(docs_path, "112823", pdf_path)
assert len(doc.sentences) == 799
# Mention Extraction
part_ngrams = MentionNgrams(n_max=1)
Part = mention_subclass("Part")
mention_extractor_udf = MentionExtractorUDF([Part], [part_ngrams],
[part_matcher])
doc = mention_extractor_udf.apply(doc)
assert doc.name == "112823"
assert len(doc.parts) == 62
part = doc.parts[0]
logger.info(f"Part: {part.context}")
# Candidate Extraction
PartRel = candidate_subclass("PartRel", [Part])
candidate_extractor_udf = CandidateExtractorUDF([PartRel], None, False,
False, True)
doc = candidate_extractor_udf.apply(doc, split=0)
# Featurization based on default feature library
featurizer_udf = FeaturizerUDF([PartRel], FeatureExtractor())
# Test that featurization default feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_default_feats = len(key_set)
# Featurization with only textual feature
feature_extractors = FeatureExtractor(features=["textual"])
featurizer_udf = FeaturizerUDF([PartRel],
feature_extractors=feature_extractors)
# Test that featurization textual feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_textual_features = len(key_set)
# Featurization with only tabular feature
feature_extractors = FeatureExtractor(features=["tabular"])
featurizer_udf = FeaturizerUDF([PartRel],
feature_extractors=feature_extractors)
# Test that featurization tabular feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_tabular_features = len(key_set)
# Featurization with only structural feature
feature_extractors = FeatureExtractor(features=["structural"])
featurizer_udf = FeaturizerUDF([PartRel],
feature_extractors=feature_extractors)
# Test that featurization structural feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_structural_features = len(key_set)
# Featurization with only visual feature
feature_extractors = FeatureExtractor(features=["visual"])
featurizer_udf = FeaturizerUDF([PartRel],
feature_extractors=feature_extractors)
# Test that featurization visual feature library
features_list = featurizer_udf.apply(doc)
features = itertools.chain.from_iterable(features_list)
key_set = set([key for feature in features for key in feature["keys"]])
n_visual_features = len(key_set)
assert (n_default_feats == n_textual_features + n_tabular_features +
n_structural_features + n_visual_features)
def test_cand_gen():
"""Test extracting candidates from mentions from documents."""
def do_nothing_matcher(fig):
return True
docs_path = "tests/data/html/112823.html"
pdf_path = "tests/data/pdf/112823.pdf"
doc = parse_doc(docs_path, "112823", pdf_path)
# Mention Extraction
part_ngrams = MentionNgramsPart(parts_by_doc=None, n_max=3)
temp_ngrams = MentionNgramsTemp(n_max=2)
volt_ngrams = MentionNgramsVolt(n_max=1)
figs = MentionFigures(types="png")
Part = mention_subclass("Part")
Temp = mention_subclass("Temp")
Volt = mention_subclass("Volt")
Fig = mention_subclass("Fig")
fig_matcher = LambdaFunctionFigureMatcher(func=do_nothing_matcher)
with pytest.raises(ValueError):
MentionExtractor(
"dummy",
[Part, Temp, Volt],
[part_ngrams, volt_ngrams], # Fail, mismatched arity
[part_matcher, temp_matcher, volt_matcher],
)
with pytest.raises(ValueError):
MentionExtractor(
"dummy",
[Part, Temp, Volt],
[part_ngrams, temp_matcher, volt_ngrams],
[part_matcher, temp_matcher], # Fail, mismatched arity
)
mention_extractor_udf = MentionExtractorUDF(
[Part, Temp, Volt, Fig],
[part_ngrams, temp_ngrams, volt_ngrams, figs],
[part_matcher, temp_matcher, volt_matcher, fig_matcher],
)
doc = mention_extractor_udf.apply(doc)
assert len(doc.parts) == 70
assert len(doc.volts) == 33
assert len(doc.temps) == 23
assert len(doc.figs) == 31
part = doc.parts[0]
volt = doc.volts[0]
temp = doc.temps[0]
logger.info(f"Part: {part.context}")
logger.info(f"Volt: {volt.context}")
logger.info(f"Temp: {temp.context}")
# Candidate Extraction
PartTemp = candidate_subclass("PartTemp", [Part, Temp])
PartVolt = candidate_subclass("PartVolt", [Part, Volt])
with pytest.raises(ValueError):
CandidateExtractor(
"dummy",
[PartTemp, PartVolt],
throttlers=[
temp_throttler,
volt_throttler,
volt_throttler,
], # Fail, mismatched arity
)
with pytest.raises(ValueError):
CandidateExtractor(
"dummy",
[PartTemp], # Fail, mismatched arity
throttlers=[temp_throttler, volt_throttler],
)
# Test that no throttler in candidate extractor
candidate_extractor_udf = CandidateExtractorUDF(
[PartTemp, PartVolt], [None, None], False, False, True # Pass, no throttler
)
doc = candidate_extractor_udf.apply(doc, split=0)
assert len(doc.part_temps) == 1610
assert len(doc.part_volts) == 2310
# Clear
doc.part_temps = []
doc.part_volts = []
# Test with None in throttlers in candidate extractor
candidate_extractor_udf = CandidateExtractorUDF(
[PartTemp, PartVolt], [temp_throttler, None], False, False, True
)
doc = candidate_extractor_udf.apply(doc, split=0)
assert len(doc.part_temps) == 1432
assert len(doc.part_volts) == 2310
# Clear
doc.part_temps = []
doc.part_volts = []
candidate_extractor_udf = CandidateExtractorUDF(
[PartTemp, PartVolt], [temp_throttler, volt_throttler], False, False, True
)
doc = candidate_extractor_udf.apply(doc, split=0)
assert len(doc.part_temps) == 1432
assert len(doc.part_volts) == 1993
assert len(doc.parts) == 70
assert len(doc.volts) == 33
assert len(doc.temps) == 23