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():