use displacy to render tags

app.py

@@ -1,8 +1,13 @@
+import pandas as pd
+from spacy import displacy
+from spacy.tokens import Doc
+from spacy.vocab import Vocab
+from spacy_streamlit.util import get_html
 import streamlit as st
 import torch
 from transformers import BertTokenizerFast
 
-from model import BertForTokenAndSequenceJointClassification
+from model import BertForTokenAndSequenceJointClassification, TOKEN_TAGS
 
 @st.cache(allow_output_mutation=True)
 def load_model():
@@ -28,4 +33,33 @@ token_class_index = torch.argmax(outputs.token_logits, dim=-1)
 tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids[0][1:-1])
 tags = [model.token_tags[i] for i in token_class_index[0].tolist()[1:-1]]
 
-st.table(list(zip(tokens, tags)))
+spaces = [not tok.startswith('##') for tok in tokens][1:] + [False]
+
+doc = Doc(Vocab(strings=set(tokens)),
+          words=tokens,
+          spaces=spaces,
+          ents=[tag if tag == "O" else f"I-{tag}" for tag in tags])
+
+labels = TOKEN_TAGS[2:]
+
+label_select = st.multiselect(
+    "Tags",
+    options=labels,
+    default=labels,
+    key=f"tags_ner_label_select",
+)
+html = displacy.render(
+    doc, style="ent", options={"ents": label_select, "colors": {}}
+)
+style = "<style>mark.entity { display: inline-block }</style>"
+st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)
+
+attrs = ["text", "label_", "start", "end", "start_char", "end_char"]
+data = [
+    [str(getattr(ent, attr)) for attr in attrs]
+    for ent in doc.ents
+    if ent.label_ in label_select
+]
+if data:
+    df = pd.DataFrame(data, columns=attrs)
+    st.dataframe(df)

requirements.txt

@@ -1,3 +1,6 @@
+pandas
+spacy
+spacy_streamlit
 streamlit
 transformers
 torch