Function test_impira_dataset

tests/test_classification_end_to_end.py:45–49

(example, model)

@pytest.mark.parametrize("example", EXAMPLES)
@pytest.mark.parametrize("model", CHECKPOINTS.keys())
def test_impira_dataset(example, model):
    document = load_document(example.path)
    pipe = pipeline("document-classification", model=CHECKPOINTS[model])
    resp = pipe(top_k=1, **document.context)
    assert resp == [{"label": x} for x in example.classes[model]]