WP

app.py

@@ -1,13 +1,11 @@
 import json
 
 import gradio as gr
+from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForTokenClassification
 
-
-def ner(text):
-    api = gr.Interface.load("d4data/biomedical-ner-all", src="models")
-    result = api(text)
-    return result
-
+tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
+model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
 
 EXAMPLE_TEXTS = []
 with open("examples.json", "r") as f:
@@ -15,10 +13,31 @@ with open("examples.json", "r") as f:
     EXAMPLE_TEXTS = [x["text"] for x in example_json]
 
 
+pipe = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
+
+
+def ner(text):
+    raw = pipe(text)
+    result = {
+        "text": text,
+        "entities": [
+            {
+                "entity": x["entity_group"],
+                "word": x["word"],
+                "score": x["score"],
+                "start": x["start"],
+                "end": x["end"],
+            }
+            for x in raw
+        ],
+    }
+    return result, {}
+
+
 interface = gr.Interface(
     ner,
     inputs=gr.Textbox(label="Input", value=""),
-    outputs=["json"],
+    outputs=[gr.HighlightedText(combine_adjacent=True), "json"],
     examples=EXAMPLE_TEXTS,
 )