update

app.py

@@ -1,36 +1,42 @@
 import gradio as gr
-from dataclasses import dataclass
-
-from pytorch_ie.annotations import LabeledSpan
-from pytorch_ie.auto import AutoPipeline
-from pytorch_ie.core import AnnotationList, annotation_field
-from pytorch_ie.documents import TextDocument
-
+from tner import TransformersNER
 from spacy import displacy
 
+model = TransformersNER("tner/roberta-large-ontonotes5")
 
-@dataclass
-class ExampleDocument(TextDocument):
-    entities: AnnotationList[LabeledSpan] = annotation_field(target="text")
-
+# DUMMY = {
+#    'prediction': [['B-person', 'I-person', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-location']],
+#    'probability': [[0.9967652559280396, 0.9994561076164246, 0.9986955523490906, 0.9947081804275513, 0.6129112243652344, 0.9984312653541565, 0.9868122935295105, 0.9983410835266113, 0.9995284080505371, 0.9838910698890686]],
+#    'input': [['Jacob', 'Collier', 'is', 'a', 'Grammy', 'awarded', 'English', 'artist', 'from', 'London']],
+#    'entity_prediction': [[
+#        {'type': 'person', 'entity': ['Jacob', 'Collier'], 'position': [0, 1], 'probability': [0.9967652559280396, 0.9994561076164246]},
+#        {'type': 'location', 'entity': ['London'], 'position': [9], 'probability': [0.9838910698890686]}
+#     ]]
+# }
 
-model_name_or_path = "pie/example-ner-spanclf-conll03"
-
-ner_pipeline = AutoPipeline.from_pretrained(model_name_or_path, device=-1, num_workers=0)
+examples = [
+    "Jacob Collier is a Grammy awarded artist from England.",
+    "When Sebastian Thrun PERSON started working on self-driving cars at Google ORG in 2007 DATE , few people outside of the company took him seriously.",
+    "But Google ORGis starting from behind. The company made a late push into hardware, and Apple ORG’s Siri, available on iPhones, and Amazon ORG’s Alexa software, which runs on its Echo and Dot devices, have clear leads in consumer adoption."
+]
 
 
 def predict(text):
-    document = ExampleDocument(text)
+    output = model.predict([text])
+    tokens = output['input'][0]
+
+    def retain_char_position(p):
+        if p == 0:
+            return 0
+        return len(' '.join(tokens[:p])) + 1
 
-    ner_pipeline(document, predict_field="entities")
-    
     doc = {
-        "text": document.text,
+        "text": text,
         "ents": [{
-            "start": entity.start,
-            "end": entity.end,
-            "label": entity.label
-        } for entity in sorted(document.entities.predictions, key=lambda e: e.start)],
+            "start": retain_char_position(entity['position'][0]),
+            "end": retain_char_position(entity['position'][-1]) + 1 + len(entity['entity'][-1]),
+            "label": entity['type']
+        } for entity in output['entity_prediction'][0]],
         "title": None
     }
 
@@ -48,8 +54,10 @@ iface = gr.Interface(
     fn=predict,
     inputs=gr.inputs.Textbox(
         lines=5,
-        default="There is still some uncertainty that Musk - also chief executive of electric car maker Tesla and rocket company SpaceX - will pull off his planned buyout.",
+        placeholder="Input Sentence",
+        default=examples[0],
+        examples=examples
     ),
     outputs="html",
 )
-iface.launch()
+iface.launch()

requirements.txt

@@ -1,2 +1,2 @@
-pytorch-ie>=0.6.0
-spacy
+spacy
+tner