Implemented NER into the Streamlit interface.

app.py

@@ -1,4 +1,3 @@
-import glob
 import random
 import os
 import numpy as np
@@ -9,6 +8,7 @@ from datasets import load_dataset
 from scipy.io.wavfile import write
 
 from modules.diarization.nemo_diarization import diarization
+from modules.nlp.nemo_ner import detect_ner
 
 FOLDER_WAV_DB = "data/database/"
 FOLDER_USER_DATA = "data/user_data/"
@@ -32,12 +32,19 @@ if st.button('Try a random sample from the database'):
     st.audio(audio_file.read())
     st.write("Starting transcription. Estimated processing time: %0.1f seconds" % (f.frames / (f.samplerate * 5)))
     result = diarization(os.path.join(FOLDER_WAV_DB, file_name + '.wav'))
+    sentences = result[file_name]["sentences"]
+    all_strings = ""
+    for sentence in sentences:
+        all_strings = all_strings + sentence["sentence"] + "\n"
     with open("info/transcripts/pred_rttms/" + file_name + ".txt") as f:
         transcript = f.read()
-    st.write("Transcription completed.")
+    st.write("Transcription completed. Starting named entity recognition.")
+    tagged_string, tags_summary = detect_ner(all_strings)
+    transcript = transcript + '\n' + tagged_string
     st.write("Number of speakers: %s" % result[file_name]["speaker_count"])
     st.write("Sentences: %s" % len(result[file_name]["sentences"]))
     st.write("Words: %s" % len(result[file_name]["words"]))
+    st.write("Found named entities: %s" % tags_summary)
     st.download_button(
         label="Download audio transcript",
         data=transcript,

modules/nlp/nemo_ner.py

@@ -9,8 +9,8 @@ pretrained_ner_model = nemo_nlp.models.TokenClassificationModel.from_pretrained(
     model_name="ner_en_bert", override_config_path=new_config)
 
 
-def detect_ner(input_string):
-    tagged_string = pretrained_ner_model.add_predictions([input_string.replace('[', '').replace(']', '')])[0]
+def detect_ner(input_strings):
+    tagged_string = pretrained_ner_model.add_predictions([input_strings.replace('[', '').replace(']', '')])[0]
     tags = re.findall('\[.*?]', tagged_string)
-    tags_summary = "Found named entities: " + str(dict(Counter(tags)))[1:-1]
+    tags_summary = str(dict(Counter(tags)))[1:-1]
     return tagged_string, tags_summary