def test_should_learn_similar_sequence(self): structured_document_train = SimpleStructuredDocument( SimplePage([ SimpleLine([ SimpleToken(TEXT_1, tag=TAG_1), SimpleToken(TEXT_1, tag=TAG_1), SimpleToken(TEXT_2, tag=TAG_2), SimpleToken(TEXT_3, tag=TAG_3) ]) ], bounding_box=PAGE_BOUNDING_BOX)) structured_document_test = SimpleStructuredDocument( SimplePage([ SimpleLine([ SimpleToken(TEXT_1, tag=TAG_1), SimpleToken(TEXT_2, tag=TAG_2), SimpleToken(TEXT_3, tag=TAG_3) ]) ], bounding_box=PAGE_BOUNDING_BOX)) token_props_list_train = list( structured_document_to_token_props(structured_document_train)) X_train = [token_props_list_to_features(token_props_list_train)] y_train = [token_props_list_to_labels(token_props_list_train)] token_props_list_test = list( structured_document_to_token_props(structured_document_test)) X_test = [token_props_list_to_features(token_props_list_test)] y_test = [token_props_list_to_labels(token_props_list_test)] with create_crf_suite_model() as model: model.fit(X_train, y_train) y_predicted = model.predict(X_test) assert y_predicted == y_test
def test_should_add_previous_and_next_token_word_features(self): result = token_props_list_to_features([ create_token_props(TEXT_1), create_token_props(TEXT_2), create_token_props(TEXT_3) ]) assert [x.get('word.lower') for x in result] == [ TEXT_1.lower(), TEXT_2.lower(), TEXT_3.lower() ] assert [x.get('-2:word.lower') for x in result] == [ None, None, TEXT_1.lower() ] assert [x.get('-1:word.lower') for x in result] == [ None, TEXT_1.lower(), TEXT_2.lower() ] assert [x.get('1:word.lower') for x in result] == [ TEXT_2.lower(), TEXT_3.lower(), None ] assert [x.get('2:word.lower') for x in result] == [ TEXT_3.lower(), None, None ] assert [x.get('BOD[-2]') for x in result] == [ True, True, None ] assert [x.get('BOD[-1]') for x in result] == [ True, None, None ] assert [x.get('EOD[1]') for x in result] == [ None, None, True ] assert [x.get('EOD[2]') for x in result] == [ None, True, True ]
def test_should_extract_scoped_tags(self): token_props = create_token_props(TEXT_1) token_props['scoped_tags'] = { SCOPE_1: TAG_1 } result = token_props_list_to_features([token_props]) assert [x.get('%s.tag' % SCOPE_1) for x in result] == [TAG_1]
def train_model(file_list, cv_file_list, cv_source_tag_scope, page_range=None, progress=True): stop_watch_recorder = StopWatchRecorder() model = CrfSuiteModel() stop_watch_recorder.start('loading files') token_props_list_by_document = load_token_props_list_by_document( file_list, cv_file_list, cv_source_tag_scope=cv_source_tag_scope, page_range=page_range, progress=progress) assert token_props_list_by_document stop_watch_recorder.start('converting to features') X = [token_props_list_to_features(x) for x in token_props_list_by_document] y = [token_props_list_to_labels(x) for x in token_props_list_by_document] get_logger().info('training model (with %d documents)', len(X)) stop_watch_recorder.start('train') model.fit(X, y) stop_watch_recorder.start('serialize') serialized_model = serialize_model(model) stop_watch_recorder.stop() get_logger().info('timings: %s', stop_watch_recorder) return serialized_model
def test_should_extract_various_word_features(self): result = token_props_list_to_features([create_token_props('TestMe')]) assert [x.get('word.lower') for x in result] == ['testme'] assert [x.get('word[:1]') for x in result] == ['t'] assert [x.get('word[-3:]') for x in result] == ['tme'] assert [x.get('word[-2:]') for x in result] == ['me'] assert [x.get('word[:1].isupper') for x in result] == [True] assert [x.get('word.isupper') for x in result] == [False] assert [x.get('word.isdigit') for x in result] == [False]
def predict_and_annotate_structured_document(structured_document, model, tag_scope=CRF_TAG_SCOPE): token_props = list(structured_document_to_token_props(structured_document)) x = token_props_list_to_features(token_props) y_pred = model.predict([x])[0] annotate_structured_document_using_predictions(structured_document, y_pred, token_props, tag_scope=tag_scope) return structured_document
def test_should_predict_and_annotate_single_token(self): token_1 = SimpleToken(TOKEN_TEXT_1, bounding_box=BOUNDING_BOX) structured_document = SimpleStructuredDocument( SimplePage(lines=[SimpleLine([token_1])], bounding_box=BOUNDING_BOX)) model = MagicMock() model.predict.return_value = [[TAG_1]] token_props = list( structured_document_to_token_props(structured_document)) X = [token_props_list_to_features(token_props)] predict_and_annotate_structured_document(structured_document, model) assert structured_document.get_tag(token_1, scope=CRF_TAG_SCOPE) == TAG_1 model.predict.assert_called_with(X)
def test_should_learn_simple_sequence(self): structured_document = SimpleStructuredDocument( SimplePage([ SimpleLine([ SimpleToken(TEXT_1, tag=TAG_1), SimpleToken(TEXT_2, tag=TAG_2), SimpleToken(TEXT_3, tag=TAG_3) ]) ], bounding_box=PAGE_BOUNDING_BOX)) token_props_list = list( structured_document_to_token_props(structured_document)) get_logger().debug('token_props_list:\n%s', token_props_list) X = [token_props_list_to_features(token_props_list)] y = [token_props_list_to_labels(token_props_list)] get_logger().debug('X:\n%s', X) get_logger().debug('y:\n%s', y) with create_crf_suite_model() as model: model.fit(X, y) y_predicted = model.predict(X) assert y_predicted == y
def test_should_pickle_and_unpickle_model(self): structured_document = SimpleStructuredDocument( SimplePage([ SimpleLine([ SimpleToken(TEXT_1, tag=TAG_1), SimpleToken(TEXT_2, tag=TAG_2), SimpleToken(TEXT_3, tag=TAG_3) ]) ], bounding_box=PAGE_BOUNDING_BOX)) token_props_list = list( structured_document_to_token_props(structured_document)) X = [token_props_list_to_features(token_props_list)] y = [token_props_list_to_labels(token_props_list)] with create_crf_suite_model() as model: model.fit(X, y) serialized_model = pickle.dumps(model) model = pickle.loads(serialized_model) y_predicted = model.predict(X) assert y_predicted == y
def test_should_not_include_tag(self): result = token_props_list_to_features([ create_token_props(TEXT_1, tag=TAG_1) ]) assert [x.get('tag') for x in result] == [None]