From c2b6c930a3ef98893c73916dab2daaa9e095731e Mon Sep 17 00:00:00 2001 From: Kasper Fyhn Date: Wed, 17 Jan 2024 08:23:51 +0100 Subject: [PATCH] fix: check for correct 'emotionally_laden' doc extension when creating emotion classifier --- src/dacy/sentiment/wrapped_models.py | 4 ++-- tests/test_sentiment.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/dacy/sentiment/wrapped_models.py b/src/dacy/sentiment/wrapped_models.py index 5e8ada9e..d4043404 100644 --- a/src/dacy/sentiment/wrapped_models.py +++ b/src/dacy/sentiment/wrapped_models.py @@ -139,7 +139,7 @@ def make_emotion_transformer( doc_extension_prediction: str, labels: List[str], # type: ignore ) -> SequenceClassificationTransformer: - if not Doc.has_extension("dacy/emotionally_laden"): + if not Doc.has_extension("emotionally_laden"): warn( "The 'emotion' component assumes the 'emotionally_laden' extension is set." + " To set it you can run nlp.add_pipe('dacy/emotionally_laden')", @@ -162,7 +162,7 @@ def make_emotion_transformer( # overwrite extension such that it return no emotion if the document does not have # an emotion - if Doc.has_extension("dacy/emotionally_laden"): + if Doc.has_extension("emotionally_laden"): def label_getter(doc) -> Optional[str]: # noqa: ANN001 # type: ignore if doc._.emotionally_laden == "emotional": diff --git a/tests/test_sentiment.py b/tests/test_sentiment.py index 7b4dd1f5..0777c92c 100644 --- a/tests/test_sentiment.py +++ b/tests/test_sentiment.py @@ -36,8 +36,10 @@ def test_add_berttone_polarity(): def test_add_bertemotion_laden(): nlp = spacy.blank("da") nlp.add_pipe("dacy/emotionally_laden") + nlp.add_pipe("dacy/emotion") doc = nlp("Der er et træ i haven") assert doc._.emotionally_laden == "no emotion" + assert doc._.emotion == "no emotion" def test_add_bertemotion_emo():