import spacy_streamlit
import streamlit as st
import typer
from scripts.torch_ner_model import build_torch_ner_model
from scripts.torch_ner_pipe import make_torch_entity_recognizer


def main(models: str = None, default_text: str = None):
    st.title('NER Predictor')
    
    models = "training/model-best,training_trf/model-best"
    default_text = "The patient had surgery."
    models = [name.strip() for name in models.split(",")]
    labels = ["person", "problem", "pronoun", "test", "treatment"]

    #if st.button('Predict entities'):
    spacy_streamlit.visualize(models, default_text, visualizers=["ner"], ner_labels=labels)


if __name__ == "__main__":
    try:
        typer.run(main)
    except SystemExit:
        pass