def load_model_and_predict(algorithm_chosen, featurizer_output):
session = featurizer_output['session']
cands = featurizer_output['candidate_variable']
featurizer = featurizer_output['featurizer_variable']
if algorithm_chosen == 'logistic_regression':
disc_model = LogisticRegression()
elif algorithm_chosen == 'sparse_logistic_regression':
disc_model = SparseLogisticRegression()
else:
disc_model = LSTM()
# Manually load settings and cardinality from a saved trained model.
checkpoint = torch.load(config.base_dir + '/checkpoints/' +
algorithm_chosen)
disc_model.settings = checkpoint["config"]
disc_model.cardinality = checkpoint["cardinality"]
# Build a model using the loaded settings and cardinality.
disc_model._build_model()
disc_model.load(model_file=algorithm_chosen,
save_dir=config.base_dir + '/checkpoints')
cand_list = [session.query(cands[0]).all()]
cand_feature_matrix = featurizer.get_feature_matrices(cand_list)
test_score = disc_model.predict((cand_list[0], cand_feature_matrix[0]),
b=0.5,
pos_label=TRUE)
true_pred = [
cand_list[0][_] for _ in np.nditer(np.where(test_score == TRUE))
]
return true_pred
def test_e2e():
"""Run an end-to-end test on documents of the hardware domain."""
PARALLEL = 4
max_docs = 12
fonduer.init_logging(
log_dir="log_folder",
format="[%(asctime)s][%(levelname)s] %(name)s:%(lineno)s - %(message)s",
level=logging.INFO,
)
session = fonduer.Meta.init(CONN_STRING).Session()
docs_path = "tests/data/html/"
pdf_path = "tests/data/pdf/"
doc_preprocessor = HTMLDocPreprocessor(docs_path, max_docs=max_docs)
corpus_parser = Parser(
session,
parallelism=PARALLEL,
structural=True,
lingual=True,
visual=True,
pdf_path=pdf_path,
)
corpus_parser.apply(doc_preprocessor)
assert session.query(Document).count() == max_docs
num_docs = session.query(Document).count()
logger.info(f"Docs: {num_docs}")
assert num_docs == max_docs
num_sentences = session.query(Sentence).count()
logger.info(f"Sentences: {num_sentences}")
# Divide into test and train
docs = sorted(corpus_parser.get_documents())
last_docs = sorted(corpus_parser.get_last_documents())
ld = len(docs)
assert ld == len(last_docs)
assert len(docs[0].sentences) == len(last_docs[0].sentences)
assert len(docs[0].sentences) == 799
assert len(docs[1].sentences) == 663
assert len(docs[2].sentences) == 784
assert len(docs[3].sentences) == 661
assert len(docs[4].sentences) == 513
assert len(docs[5].sentences) == 700
assert len(docs[6].sentences) == 528
assert len(docs[7].sentences) == 161
assert len(docs[8].sentences) == 228
assert len(docs[9].sentences) == 511
assert len(docs[10].sentences) == 331
assert len(docs[11].sentences) == 528
# Check table numbers
assert len(docs[0].tables) == 9
assert len(docs[1].tables) == 9
assert len(docs[2].tables) == 14
assert len(docs[3].tables) == 11
assert len(docs[4].tables) == 11
assert len(docs[5].tables) == 10
assert len(docs[6].tables) == 10
assert len(docs[7].tables) == 2
assert len(docs[8].tables) == 7
assert len(docs[9].tables) == 10
assert len(docs[10].tables) == 6
assert len(docs[11].tables) == 9
# Check figure numbers
assert len(docs[0].figures) == 32
assert len(docs[1].figures) == 11
assert len(docs[2].figures) == 38
assert len(docs[3].figures) == 31
assert len(docs[4].figures) == 7
assert len(docs[5].figures) == 38
assert len(docs[6].figures) == 10
assert len(docs[7].figures) == 31
assert len(docs[8].figures) == 4
assert len(docs[9].figures) == 27
assert len(docs[10].figures) == 5
assert len(docs[11].figures) == 27
# Check caption numbers
assert len(docs[0].captions) == 0
assert len(docs[1].captions) == 0
assert len(docs[2].captions) == 0
assert len(docs[3].captions) == 0
assert len(docs[4].captions) == 0
assert len(docs[5].captions) == 0
assert len(docs[6].captions) == 0
assert len(docs[7].captions) == 0
assert len(docs[8].captions) == 0
assert len(docs[9].captions) == 0
assert len(docs[10].captions) == 0
assert len(docs[11].captions) == 0
train_docs = set()
dev_docs = set()
test_docs = set()
splits = (0.5, 0.75)
data = [(doc.name, doc) for doc in docs]
data.sort(key=lambda x: x[0])
for i, (doc_name, doc) in enumerate(data):
if i < splits[0] * ld:
train_docs.add(doc)
elif i < splits[1] * ld:
dev_docs.add(doc)
else:
test_docs.add(doc)
logger.info([x.name for x in train_docs])
# NOTE: With multi-relation support, return values of getting candidates,
# mentions, or sparse matrices are formatted as a list of lists. This means
# that with a single relation, we need to index into the list of lists to
# get the candidates/mentions/sparse matrix for a particular relation or
# mention.
# Mention Extraction
part_ngrams = MentionNgramsPart(parts_by_doc=None, n_max=3)
temp_ngrams = MentionNgramsTemp(n_max=2)
volt_ngrams = MentionNgramsVolt(n_max=1)
Part = mention_subclass("Part")
Temp = mention_subclass("Temp")
Volt = mention_subclass("Volt")
mention_extractor = MentionExtractor(
session,
[Part, Temp, Volt],
[part_ngrams, temp_ngrams, volt_ngrams],
[part_matcher, temp_matcher, volt_matcher],
)
mention_extractor.apply(docs, parallelism=PARALLEL)
assert session.query(Part).count() == 299
assert session.query(Temp).count() == 138
assert session.query(Volt).count() == 140
assert len(mention_extractor.get_mentions()) == 3
assert len(mention_extractor.get_mentions()[0]) == 299
assert (
len(
mention_extractor.get_mentions(
docs=[session.query(Document).filter(Document.name == "112823").first()]
)[0]
)
== 70
)
# Candidate Extraction
PartTemp = candidate_subclass("PartTemp", [Part, Temp])
PartVolt = candidate_subclass("PartVolt", [Part, Volt])
candidate_extractor = CandidateExtractor(
session, [PartTemp, PartVolt], throttlers=[temp_throttler, volt_throttler]
)
for i, docs in enumerate([train_docs, dev_docs, test_docs]):
candidate_extractor.apply(docs, split=i, parallelism=PARALLEL)
assert session.query(PartTemp).filter(PartTemp.split == 0).count() == 3493
assert session.query(PartTemp).filter(PartTemp.split == 1).count() == 61
assert session.query(PartTemp).filter(PartTemp.split == 2).count() == 416
assert session.query(PartVolt).count() == 4282
# Grab candidate lists
train_cands = candidate_extractor.get_candidates(split=0, sort=True)
dev_cands = candidate_extractor.get_candidates(split=1, sort=True)
test_cands = candidate_extractor.get_candidates(split=2, sort=True)
assert len(train_cands) == 2
assert len(train_cands[0]) == 3493
assert (
len(
candidate_extractor.get_candidates(
docs=[session.query(Document).filter(Document.name == "112823").first()]
)[0]
)
== 1432
)
# Featurization
featurizer = Featurizer(session, [PartTemp, PartVolt])
# Test that FeatureKey is properly reset
featurizer.apply(split=1, train=True, parallelism=PARALLEL)
assert session.query(Feature).count() == 214
assert session.query(FeatureKey).count() == 1260
# Test Dropping FeatureKey
# Should force a row deletion
featurizer.drop_keys(["DDL_e1_W_LEFT_POS_3_[NNP NN IN]"])
assert session.query(FeatureKey).count() == 1259
# Should only remove the part_volt as a relation and leave part_temp
assert set(
session.query(FeatureKey)
.filter(FeatureKey.name == "DDL_e1_LEMMA_SEQ_[bc182]")
.one()
.candidate_classes
) == {"part_temp", "part_volt"}
featurizer.drop_keys(["DDL_e1_LEMMA_SEQ_[bc182]"], candidate_classes=[PartVolt])
assert session.query(FeatureKey).filter(
FeatureKey.name == "DDL_e1_LEMMA_SEQ_[bc182]"
).one().candidate_classes == ["part_temp"]
assert session.query(FeatureKey).count() == 1259
# Inserting the removed key
featurizer.upsert_keys(
["DDL_e1_LEMMA_SEQ_[bc182]"], candidate_classes=[PartTemp, PartVolt]
)
assert set(
session.query(FeatureKey)
.filter(FeatureKey.name == "DDL_e1_LEMMA_SEQ_[bc182]")
.one()
.candidate_classes
) == {"part_temp", "part_volt"}
assert session.query(FeatureKey).count() == 1259
# Removing the key again
featurizer.drop_keys(["DDL_e1_LEMMA_SEQ_[bc182]"], candidate_classes=[PartVolt])
# Removing the last relation from a key should delete the row
featurizer.drop_keys(["DDL_e1_LEMMA_SEQ_[bc182]"], candidate_classes=[PartTemp])
assert session.query(FeatureKey).count() == 1258
session.query(Feature).delete(synchronize_session="fetch")
session.query(FeatureKey).delete(synchronize_session="fetch")
featurizer.apply(split=0, train=True, parallelism=PARALLEL)
assert session.query(Feature).count() == 6478
assert session.query(FeatureKey).count() == 4538
F_train = featurizer.get_feature_matrices(train_cands)
assert F_train[0].shape == (3493, 4538)
assert F_train[1].shape == (2985, 4538)
assert len(featurizer.get_keys()) == 4538
featurizer.apply(split=1, parallelism=PARALLEL)
assert session.query(Feature).count() == 6692
assert session.query(FeatureKey).count() == 4538
F_dev = featurizer.get_feature_matrices(dev_cands)
assert F_dev[0].shape == (61, 4538)
assert F_dev[1].shape == (153, 4538)
featurizer.apply(split=2, parallelism=PARALLEL)
assert session.query(Feature).count() == 8252
assert session.query(FeatureKey).count() == 4538
F_test = featurizer.get_feature_matrices(test_cands)
assert F_test[0].shape == (416, 4538)
assert F_test[1].shape == (1144, 4538)
gold_file = "tests/data/hardware_tutorial_gold.csv"
labeler = Labeler(session, [PartTemp, PartVolt])
labeler.apply(
docs=last_docs,
lfs=[[gold], [gold]],
table=GoldLabel,
train=True,
parallelism=PARALLEL,
)
assert session.query(GoldLabel).count() == 8252
stg_temp_lfs = [
LF_storage_row,
LF_operating_row,
LF_temperature_row,
LF_tstg_row,
LF_to_left,
LF_negative_number_left,
]
ce_v_max_lfs = [
LF_bad_keywords_in_row,
LF_current_in_row,
LF_non_ce_voltages_in_row,
]
with pytest.raises(ValueError):
labeler.apply(split=0, lfs=stg_temp_lfs, train=True, parallelism=PARALLEL)
labeler.apply(
docs=train_docs,
lfs=[stg_temp_lfs, ce_v_max_lfs],
train=True,
parallelism=PARALLEL,
)
assert session.query(Label).count() == 6478
assert session.query(LabelKey).count() == 9
L_train = labeler.get_label_matrices(train_cands)
assert L_train[0].shape == (3493, 9)
assert L_train[1].shape == (2985, 9)
assert len(labeler.get_keys()) == 9
# Test Dropping LabelerKey
labeler.drop_keys(["LF_storage_row"])
assert len(labeler.get_keys()) == 8
# Test Upserting LabelerKey
labeler.upsert_keys(["LF_storage_row"])
assert "LF_storage_row" in [label.name for label in labeler.get_keys()]
L_train_gold = labeler.get_gold_labels(train_cands)
assert L_train_gold[0].shape == (3493, 1)
L_train_gold = labeler.get_gold_labels(train_cands, annotator="gold")
assert L_train_gold[0].shape == (3493, 1)
gen_model = LabelModel()
gen_model.fit(L_train=L_train[0], n_epochs=500, log_freq=100)
train_marginals = gen_model.predict_proba(L_train[0])
disc_model = LogisticRegression()
disc_model.train(
(train_cands[0], F_train[0]),
train_marginals,
X_dev=(train_cands[0], F_train[0]),
Y_dev=L_train_gold[0].reshape(-1),
b=0.6,
pos_label=TRUE,
n_epochs=5,
lr=0.001,
)
test_score = disc_model.predict((test_cands[0], F_test[0]), b=0.6, pos_label=TRUE)
true_pred = [test_cands[0][_] for _ in np.nditer(np.where(test_score == TRUE))]
pickle_file = "tests/data/parts_by_doc_dict.pkl"
with open(pickle_file, "rb") as f:
parts_by_doc = pickle.load(f)
(TP, FP, FN) = entity_level_f1(
true_pred, gold_file, ATTRIBUTE, test_docs, parts_by_doc=parts_by_doc
)
tp_len = len(TP)
fp_len = len(FP)
fn_len = len(FN)
prec = tp_len / (tp_len + fp_len) if tp_len + fp_len > 0 else float("nan")
rec = tp_len / (tp_len + fn_len) if tp_len + fn_len > 0 else float("nan")
f1 = 2 * (prec * rec) / (prec + rec) if prec + rec > 0 else float("nan")
logger.info(f"prec: {prec}")
logger.info(f"rec: {rec}")
logger.info(f"f1: {f1}")
assert f1 < 0.7 and f1 > 0.3
stg_temp_lfs_2 = [
LF_to_left,
LF_test_condition_aligned,
LF_collector_aligned,
LF_current_aligned,
LF_voltage_row_temp,
LF_voltage_row_part,
LF_typ_row,
LF_complement_left_row,
LF_too_many_numbers_row,
LF_temp_on_high_page_num,
LF_temp_outside_table,
LF_not_temp_relevant,
]
labeler.update(split=0, lfs=[stg_temp_lfs_2, ce_v_max_lfs], parallelism=PARALLEL)
assert session.query(Label).count() == 6478
assert session.query(LabelKey).count() == 16
L_train = labeler.get_label_matrices(train_cands)
assert L_train[0].shape == (3493, 16)
gen_model = LabelModel()
gen_model.fit(L_train=L_train[0], n_epochs=500, log_freq=100)
train_marginals = gen_model.predict_proba(L_train[0])
disc_model = LogisticRegression()
disc_model.train(
(train_cands[0], F_train[0]), train_marginals, n_epochs=5, lr=0.001
)
test_score = disc_model.predict((test_cands[0], F_test[0]), b=0.6, pos_label=TRUE)
true_pred = [test_cands[0][_] for _ in np.nditer(np.where(test_score == TRUE))]
(TP, FP, FN) = entity_level_f1(
true_pred, gold_file, ATTRIBUTE, test_docs, parts_by_doc=parts_by_doc
)
tp_len = len(TP)
fp_len = len(FP)
fn_len = len(FN)
prec = tp_len / (tp_len + fp_len) if tp_len + fp_len > 0 else float("nan")
rec = tp_len / (tp_len + fn_len) if tp_len + fn_len > 0 else float("nan")
f1 = 2 * (prec * rec) / (prec + rec) if prec + rec > 0 else float("nan")
logger.info(f"prec: {prec}")
logger.info(f"rec: {rec}")
logger.info(f"f1: {f1}")
assert f1 > 0.7
# Testing LSTM
disc_model = LSTM()
disc_model.train(
(train_cands[0], F_train[0]), train_marginals, n_epochs=5, lr=0.001
)
test_score = disc_model.predict((test_cands[0], F_test[0]), b=0.6, pos_label=TRUE)
true_pred = [test_cands[0][_] for _ in np.nditer(np.where(test_score == TRUE))]
(TP, FP, FN) = entity_level_f1(
true_pred, gold_file, ATTRIBUTE, test_docs, parts_by_doc=parts_by_doc
)
tp_len = len(TP)
fp_len = len(FP)
fn_len = len(FN)
prec = tp_len / (tp_len + fp_len) if tp_len + fp_len > 0 else float("nan")
rec = tp_len / (tp_len + fn_len) if tp_len + fn_len > 0 else float("nan")
f1 = 2 * (prec * rec) / (prec + rec) if prec + rec > 0 else float("nan")
logger.info(f"prec: {prec}")
logger.info(f"rec: {rec}")
logger.info(f"f1: {f1}")
assert f1 > 0.7
# Testing Sparse Logistic Regression
disc_model = SparseLogisticRegression()
disc_model.train(
(train_cands[0], F_train[0]), train_marginals, n_epochs=5, lr=0.001
)
test_score = disc_model.predict((test_cands[0], F_test[0]), b=0.6, pos_label=TRUE)
true_pred = [test_cands[0][_] for _ in np.nditer(np.where(test_score == TRUE))]
(TP, FP, FN) = entity_level_f1(
true_pred, gold_file, ATTRIBUTE, test_docs, parts_by_doc=parts_by_doc
)
tp_len = len(TP)
fp_len = len(FP)
fn_len = len(FN)
prec = tp_len / (tp_len + fp_len) if tp_len + fp_len > 0 else float("nan")
rec = tp_len / (tp_len + fn_len) if tp_len + fn_len > 0 else float("nan")
f1 = 2 * (prec * rec) / (prec + rec) if prec + rec > 0 else float("nan")
logger.info(f"prec: {prec}")
logger.info(f"rec: {rec}")
logger.info(f"f1: {f1}")
assert f1 > 0.7
# Testing Sparse LSTM
disc_model = SparseLSTM()
disc_model.train(
(train_cands[0], F_train[0]), train_marginals, n_epochs=5, lr=0.001
)
test_score = disc_model.predict((test_cands[0], F_test[0]), b=0.6, pos_label=TRUE)
true_pred = [test_cands[0][_] for _ in np.nditer(np.where(test_score == TRUE))]
(TP, FP, FN) = entity_level_f1(
true_pred, gold_file, ATTRIBUTE, test_docs, parts_by_doc=parts_by_doc
)
tp_len = len(TP)
fp_len = len(FP)
fn_len = len(FN)
prec = tp_len / (tp_len + fp_len) if tp_len + fp_len > 0 else float("nan")
rec = tp_len / (tp_len + fn_len) if tp_len + fn_len > 0 else float("nan")
f1 = 2 * (prec * rec) / (prec + rec) if prec + rec > 0 else float("nan")
logger.info(f"prec: {prec}")
logger.info(f"rec: {rec}")
logger.info(f"f1: {f1}")
assert f1 > 0.7
# Evaluate mention level scores
L_test_gold = labeler.get_gold_labels(test_cands, annotator="gold")
Y_test = L_test_gold[0].reshape(-1)
scores = disc_model.score((test_cands[0], F_test[0]), Y_test, b=0.6, pos_label=TRUE)
logger.info(scores)
assert scores["f1"] > 0.6
