def test_token_classification_viz_on_true_classes_value_error():
explainer_string = "We visited Paris during the weekend"
ner_explainer = TokenClassificationExplainer(DISTILBERT_MODEL,
DISTILBERT_TOKENIZER)
ner_explainer(explainer_string)
true_classes = ["None", "Location", "None"]
with pytest.raises(ValueError):
ner_explainer.visualize(true_classes=true_classes)
def test_token_classification_viz():
explainer_string = "We visited Paris during the weekend"
ner_explainer = TokenClassificationExplainer(DISTILBERT_MODEL,
DISTILBERT_TOKENIZER)
ner_explainer(explainer_string)
ner_explainer.visualize()