Skip to content

Commit

Permalink
fix: check for correct 'emotionally_laden' doc extension when creatin…
Browse files Browse the repository at this point in the history
…g emotion classifier
  • Loading branch information
KasperFyhn committed Jan 17, 2024
1 parent 017e13d commit c2b6c93
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/dacy/sentiment/wrapped_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')",
Expand All @@ -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":
Expand Down
2 changes: 2 additions & 0 deletions tests/test_sentiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down

0 comments on commit c2b6c93

Please sign in to comment.