Pytorch_clinical_NER / scripts /visualize_model.py
AmitGarage's picture
Upload 8 files
3ab8bd6
import spacy_streamlit
import typer
from torch_ner_model import build_torch_ner_model
from torch_ner_pipe import make_torch_entity_recognizer
def main(models: str, default_text: str):
models = [name.strip() for name in models.split(",")]
labels = ["person", "problem", "pronoun", "test", "treatment"]
spacy_streamlit.visualize(
models, default_text, visualizers=["ner"], ner_labels=labels
)
if __name__ == "__main__":
try:
typer.run(main)
except SystemExit:
pass