linting

explosion · Jan 28, 2024 · f5268c2 · f5268c2
1 parent a1d0e80
commit f5268c2
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 deletions.
diff --git a/spacy_llm/tasks/entity_linker/task.py b/spacy_llm/tasks/entity_linker/task.py
@@ -334,7 +334,11 @@ def unhighlight_ents_in_doc(doc: Doc) -> Doc:
             for ent in doc.ents
             if ent.start - 1 > 0 and doc[ent.start - 1].text == "*"
         }
-        highlight_end_idx = {ent.end for ent in doc.ents if ent.end < len(doc) and doc[ent.end].text == "*"}
+        highlight_end_idx = {
+            ent.end
+            for ent in doc.ents
+            if ent.end < len(doc) and doc[ent.end].text == "*"
+        }
         highlight_idx = highlight_start_idx | highlight_end_idx
 
         # Compute entity indices with removed highlights.

diff --git a/spacy_llm/tests/tasks/test_entity_linker.py b/spacy_llm/tests/tasks/test_entity_linker.py
@@ -683,12 +683,33 @@ def test_ent_highlighting():
         == "Alice goes to *Boston* to see the *Boston Celtics* game."
     )
 
-@pytest.mark.parametrize("text,ents,include_ents", [("Alice goes to Boston to see the Boston Celtics game.",[{"start":3, "end":4, "label":"LOC"},{"start":7, "end":9, "label":"ORG"}],[True,True]),("I went to see Boston in concert yesterday",[{"start":4, "end":5, "label":"GPE"},{"start":7, "end":8,"label":"DATE"}],[True,False])])
+
+@pytest.mark.parametrize(
+    "text,ents,include_ents",
+    [
+        (
+            "Alice goes to Boston to see the Boston Celtics game.",
+            [
+                {"start": 3, "end": 4, "label": "LOC"},
+                {"start": 7, "end": 9, "label": "ORG"},
+            ],
+            [True, True],
+        ),
+        (
+            "I went to see Boston in concert yesterday",
+            [
+                {"start": 4, "end": 5, "label": "GPE"},
+                {"start": 7, "end": 8, "label": "DATE"},
+            ],
+            [True, False],
+        ),
+    ],
+)
 def test_ent_unhighlighting(text, ents, include_ents):
     """Tests unhighlighting of entities in text."""
     nlp = spacy.blank("en")
     doc = nlp.make_doc(text)
-    doc.ents = [Span(doc=doc,**ents[0]), Span(doc=doc,**ents[1])]
+    doc.ents = [Span(doc=doc, **ents[0]), Span(doc=doc, **ents[1])]
 
     assert (
         EntityLinkerTask.unhighlight_ents_in_doc(