Upload streamlit_app.py

src/streamlit_app.py

@@ -20,6 +20,18 @@ LABEL_COLORS = {
     'LABEL-8': '#ffc6ff',  # I-SET
 }
 
+LABEL_MEANINGS = {
+    'LABEL-0': 'NONE',
+    'LABEL-1': 'B-DATE',
+    'LABEL-2': 'I-DATE',
+    'LABEL-3': 'B-TIME',
+    'LABEL-4': 'I-TIME',
+    'LABEL-5': 'B-DURATION',
+    'LABEL-6': 'I-DURATION',
+    'LABEL-7': 'B-SET',
+    'LABEL-8': 'I-SET',
+}
+
 @st.cache_resource(show_spinner=True)
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained('asdc/Bio-RoBERTime')
@@ -75,12 +87,59 @@ def colorize_entities(ner_result: List[Tuple[str, str]]) -> str:
             html += f'{token} '
     return html
 
+def extract_entities(ner_result: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    # Group consecutive tokens with the same entity label (not LABEL-0)
+    entities = []
+    current_entity = []
+    current_label = None
+    for token, label in ner_result:
+        if label != 'LABEL-0':
+            if current_label == label:
+                current_entity.append(token)
+            else:
+                if current_entity:
+                    entities.append((' '.join(current_entity), current_label))
+                current_entity = [token]
+                current_label = label
+        else:
+            if current_entity:
+                entities.append((' '.join(current_entity), current_label))
+                current_entity = []
+                current_label = None
+    if current_entity:
+        entities.append((' '.join(current_entity), current_label))
+    return entities
+
+def legend_html() -> str:
+    html = '<div style="display:flex;flex-wrap:wrap;gap:8px;">'
+    for label, color in LABEL_COLORS.items():
+        if label == 'LABEL-0':
+            continue
+        meaning = LABEL_MEANINGS[label]
+        html += f'<span style="background-color:{color};padding:2px 8px;border-radius:4px;">{meaning} ({label})</span>'
+    html += '</div>'
+    return html
+
 st.title('LLM-powered Named Entity Recognition (NER)')
 
+st.markdown('**Legend:**')
+st.markdown(legend_html(), unsafe_allow_html=True)
+
 user_text = st.text_area('Enter text for NER:', height=150)
 
 if user_text:
     ner_result = ner_with_robertime(user_text)
-    st.markdown('#### Entities:')
-    st.markdown(colorize_entities(ner_result), unsafe_allow_html=True)
+    has_entity = any(label != 'LABEL-0' for _, label in ner_result)
+    if has_entity:
+        st.markdown('#### Entities Highlighted:')
+        st.markdown(colorize_entities(ner_result), unsafe_allow_html=True)
+        entities = extract_entities(ner_result)
+        if entities:
+            st.markdown('#### Detected Entities:')
+            for ent, label in entities:
+                st.markdown(f'- <span style="background-color:{LABEL_COLORS[label]};padding:2px 8px;border-radius:4px;">{ent}</span> <span style="color:#888;">({LABEL_MEANINGS[label]})</span>', unsafe_allow_html=True)
+        else:
+            st.info('No entities detected.')
+    else:
+        st.info('No entities detected.')
     st.caption('Model: [asdc/Bio-RoBERTime](https://huggingface.co/asdc/Bio-RoBERTime)')