import spacy
import spacy_transformers
import gradio as gr

nlp = spacy.load("en_core_web_sm")

examples = [
    "Does Chicago have any stores and does Joe live here?",
]

def ner(text):
    doc = nlp(text)
    final_output = []
    flagged_categories = ["CARDINAL", "DATE", "MONEY", "PERCENT", "QUANTITY", "TIME", "ORDINAL"]
    
    for ent in doc.ents:
      label = ent.label_
  
      if label not in flagged_categories:
          output = {'entity': ent.label_, "word": ent.text, "start": int(ent.start_char), "end": int(ent.end_char)}
          final_output.append(output)
        
    return {"text": text, "entities": final_output}    

demo = gr.Interface(ner,
             gr.Textbox(placeholder="Enter sentence here..."), 
             gr.HighlightedText(),
             examples=examples)


if __name__ == "__main__":
  demo.launch(debug=True)