Updated app.py with NER model

app.py

@@ -9,6 +9,8 @@ from nltk.corpus import stopwords
 nltk.download('stopwords')
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.feature_extraction.text import TfidfVectorizer
+from flair.data import Sentence
+from flair.models import SequenceTagger
 
 # file name
 #lr_filename = 'lr_021223.pkl'
@@ -55,13 +57,35 @@ Prediction function
     #return prediction
 '''
 
+# Specify NER model
+tagger = SequenceTagger.load('/Users/julia/Models/ner_04_results/best-model.pt') # SequenceTagger.load('best-model.pt')
+
+# Runs NER on input text
+def run_ner(input_text):
+    sentence = Sentence(input_text)
+    tagger.predict(sentence)
+    output = []
+    for entity in sentence.get_spans('ner'):
+        output.append({'entity': entity.get_label('ner').value, 'word': entity.text, 'start': entity.start_position, 'end': entity.end_position})
+    return {"text": input_text, "entities": output}
+
+# Run both models, and return a tuple of their results
+def run_models(input_text):
+    prediction = 0 # This "0" is a placeholder to avoid errors; once the LR model is working, use this instead: prediction = predict(input_text)
+    entities = run_ner(input_text)
+    return prediction, entities
+
 # Define interface
-demo = gr.Interface(fn=predict,
-                        title="Text Classification Demo",
-                        description="This is a demo of a text classification model using Logistic Regression.",
-                        inputs=gr.Textbox(lines=10, placeholder='Input text here...', label="Input Text"),
-                        outputs=gr.Textbox(label="Predicted Label: Other: 0, Healthcare: 1, Technology: 2", lines=2, placeholder='Predicted label will appear here...'),
-                        allow_flagging='never'
+demo = gr.Interface(fn=run_models,
+                    title="Text Classification & Named Entity Recognition Demo",
+                    description="This is a demo of a text classification model using logistic regression as well as a named entity recognition model. Enter in some text or use one of the provided examples. Note that common named entity recognition tags include **geo** (geographical entity), **org** (organization), **per** (person), and **tim** (time).",
+                    article='*This demo is based on Logistic Regression and Named Entity Recognition models trained by Curtis Pond and Julia Nickerson as part of their FourthBrain capstone project. For more information, check out their [GitHub repo](https://github.com/nickersonj/glg-capstone).*',
+                    inputs=gr.Textbox(lines=10, placeholder='Input text here...', label="Input Text"),
+                    outputs=[gr.Textbox(label="Predicted Classification Label: Other: 0, Healthcare: 1, Technology: 2", lines=2, placeholder='Predicted label will appear here...'), 
+                             gr.HighlightedText(label='Named Entity Recognition Results')],
+                    # These examples are just placeholders; once the LR model is working, we can use longer example text such as paragraphs
+                    examples=['The indictments were announced Tuesday by the Justice Department in Cairo.', "In 2019, the men's singles winner was Novak Djokovic who defeated Roger Federer in a tournament taking place in the United Kingdom.", 'In a study published by the American Heart Association on January 18, researchers at the Johns Hopkins School of Medicine found that meal timing did not impact weight.'],
+                    allow_flagging='never'
 )
 
 demo.launch()