Update app.py

app.py

@@ -21,12 +21,16 @@ class ExampleDocument(TextDocument):
 
 ner_model_name_or_path = "pie/example-ner-spanclf-conll03"
 re_model_name_or_path = "DFKI-SLT/relation_classification_tacred_revisited"
-#"pie/example-re-textclf-tacred"
-#"DFKI-SLT/relation_classification_tacred_revisited"
 
 ner_pipeline = AutoPipeline.from_pretrained(ner_model_name_or_path, device=-1, num_workers=0)
 re_pipeline = AutoPipeline.from_pretrained(re_model_name_or_path, device=-1, num_workers=0)
 
+ner_tag_mapping = {
+    'ORG': 'ORGANIZATION',
+    'PER': 'PERSON',
+    'LOC': 'LOCATION'
+}
+
 
 def predict(text):
     document = ExampleDocument(text)
@@ -34,10 +38,16 @@ def predict(text):
     ner_pipeline(document)
 
     while len(document.entities.predictions) > 0:
-        document.entities.append(document.entities.predictions.pop(0))
+        entity = document.entities.predictions.pop(0)
+        if entity.label in ner_tag_mapping:
+            entity = LabeledSpan(start=entity.start, end=entity.end, label=ner_tag_mapping[entity.label],
+                                 score=entity.score)
+        if entity.label in re_pipeline.taskmodule.entity_labels:
+            document.entities.append(entity)
 
     re_pipeline(document)
 
+
     t = PrettyTable()
     t.field_names = ["head", "tail", "relation"]
     t.align = "l"