import spacy
import spacy_transformers
import gradio as gr

nlp = spacy.load("en_core_web_trf")

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)