fixed bug in extracting correct ner feature provider for resolution viz
@@ -146,9 +146,13 @@ class VizUtilsHC():
"""Finds entities_col,resolution_col,doc_col cols for resolution viz viz"""
entities_component, resolution_component, doc_component = None, None, None
for c in pipe.components:
- if c.name == NLP_NODE_IDS.DOCUMENT_ASSEMBLER: doc_component = c
- if c.name in [NLP_HC_NODE_IDS.NER_CONVERTER_INTERNAL, NLP_NODE_IDS.NER_CONVERTER]: entities_component = c
- if c.name == NLP_HC_NODE_IDS.SENTENCE_ENTITY_RESOLVER: resolution_component = c
+ if c.name == NLP_NODE_IDS.DOCUMENT_ASSEMBLER:
+ doc_component = c
+ if NLP_FEATURES.NAMED_ENTITY_CONVERTED in c.out_types:
+ entities_component = c
+
+ if c.name == NLP_HC_NODE_IDS.SENTENCE_ENTITY_RESOLVER:
+ resolution_component = c
entities_col = entities_component.spark_output_column_names[0]
resolution_col = resolution_component.spark_output_column_names[0]
doc_col = doc_component.spark_output_column_names[0]
@@ -198,8 +202,6 @@ class VizUtilsHC():
if user_assertion_col:
assertion_col = assertion_col
-
-
if len(viz_colors) > 0: assertion_vis.set_label_colors(viz_colors)
if write_to_streamlit:
import streamlit as st
