diff --git a/examples/03_visualize-ner-manual.py b/examples/03_visualize-ner-manual.py new file mode 100644 index 0000000..61f4b69 --- /dev/null +++ b/examples/03_visualize-ner-manual.py @@ -0,0 +1,21 @@ +""" +Example of using manual=True for visualize_ner. +""" +import spacy_streamlit +import streamlit as st + +st.title("My cool app") + +doc = [{ + "text": "But Google is starting from behind.", + "ents": [{"start": 4, "end": 10, "label": "ORG"}], + "title": None +}] + +spacy_streamlit.visualize_ner( + doc, + labels=["ORG"], + show_table=False, + title="Persons, dates and locations", + manual=True +) diff --git a/spacy_streamlit/visualizer.py b/spacy_streamlit/visualizer.py index 59b3795..d5fea4d 100644 --- a/spacy_streamlit/visualizer.py +++ b/spacy_streamlit/visualizer.py @@ -162,7 +162,7 @@ def visualize_parser( def visualize_ner( - doc: spacy.tokens.Doc, + doc: Union[spacy.tokens.Doc, List[Dict[str, str]]], *, labels: Sequence[str] = tuple(), attrs: List[str] = NER_ATTRS, @@ -170,13 +170,20 @@ def visualize_ner( title: Optional[str] = "Named Entities", colors: Dict[str, str] = {}, key: Optional[str] = None, + manual: Optional[bool] = False, ) -> None: """Visualizer for named entities.""" if title: st.header(title) - labels = labels or [ent.label_ for ent in doc.ents] - + if manual: + if show_table: + st.warning("When the parameter 'manual' is set to True, the parameter 'show_table' must be set to False.") + if not isinstance(doc, list): + st.warning("When the parameter 'manual' is set to True, the parameter 'doc' must be of type 'list', not 'spacy.tokens.Doc'.") + else: + labels = labels or [ent.label_ for ent in doc.ents] + if not labels: st.warning("The parameter 'labels' should not be empty or None.") else: @@ -188,7 +195,7 @@ def visualize_ner( key=f"{key}_ner_label_select", ) html = displacy.render( - doc, style="ent", options={"ents": label_select, "colors": colors} + doc, style="ent", options={"ents": label_select, "colors": colors}, manual=manual ) style = "" st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)