Update app.py

app.py

@@ -30,7 +30,30 @@ def run_pipeline(text):
     
     dense_output = dense_qa_pipeline(text)
     for item in dict(dense_output)['predictions'][0]:
-        dense_entities.append(dict(item))
+        dictionary = dict(item)
+        entity = dictionary['entity']
+        
+        if entity == "LABEL_0":
+            value = "O"
+        elif entity == "LABEL_1":
+            value = "B-PER"
+        elif entity == "LABEL_2":
+            value = "I-PER"
+        elif entity == "LABEL_3":
+            value = "-ORG"
+        elif entity == "LABEL_4":
+            value = "I-ORG"
+        elif entity == "LABEL_5":
+            value = "B-LOC"
+        elif entity == "LABEL_6":
+            value = "I-LOC"
+        elif entity == "LABEL_7":
+            value = "B-MISC"
+        else:
+            value = "I-MISC"
+        dictionary['entity'] = value
+        dense_entities.append(dictionary)
+        
     dense_output = {"text": text, "entities": dense_entities}
         
     dense_end = time.perf_counter()
@@ -40,7 +63,29 @@ def run_pipeline(text):
     
     sparse_output = sparse_qa_pipeline(text)
     for item in dict(sparse_output)['predictions'][0]:
-        sparse_entities.append(dict(item))
+        sparse_dictionary = dict(item)
+        entity = dictionary['entity']
+        
+        if entity == "LABEL_0":
+            value = "O"
+        elif entity == "LABEL_1":
+            value = "B-PER"
+        elif entity == "LABEL_2":
+            value = "I-PER"
+        elif entity == "LABEL_3":
+            value = "-ORG"
+        elif entity == "LABEL_4":
+            value = "I-ORG"
+        elif entity == "LABEL_5":
+            value = "B-LOC"
+        elif entity == "LABEL_6":
+            value = "I-LOC"
+        elif entity == "LABEL_7":
+            value = "B-MISC"
+        else:
+            value = "I-MISC"
+        dictionary['entity'] = value
+        sparse_entities.append(sparse_dictionary)
 
     sparse_output = {"text": text, "entities": sparse_entities}