Update app.py

app.py

@@ -5,6 +5,36 @@ import gradio as gr
 jp_model = GLiNER.from_pretrained("vumichien/ner-jp-gliner")
 meal_model = GLiNER.from_pretrained("urchade/gliner_mediumv2.1")
 
+
+def merge_tokens(entities, text):
+    # Remove spaces from the text
+    merged_text = text.replace(" ", "")
+
+    updated_entities = []
+    for entity in entities:
+        # Calculate the new start and end positions
+        start = entity['start']
+        end = entity['end']
+
+        # Get the text without spaces
+        entity_text = entity['text'].replace(" ", "")
+
+        # Find the new start and end in the merged text
+        new_start = merged_text.find(entity_text)
+        new_end = new_start + len(entity_text)
+
+        # Update the entity with new positions
+        updated_entities.append({
+            'start': new_start,
+            'end': new_end,
+            'text': entity_text,
+            'label': entity['label'],
+            'score': entity['score']
+        })
+
+    return updated_entities
+
+
 examples = [
     [
         "ner_jp",
@@ -29,23 +59,39 @@ def ner(
     labels = labels.split(",")
     if models == "ner_jp":
         model = jp_model
+        tokenized_text = " ".join(list(text))
+        entities = model.predict_entities(tokenized_text, labels, flat_ner=not nested_ner, threshold=threshold)
+        updated_entities = merge_tokens(entities, tokenized_text)
+        return {
+            "text": text,
+            "entities": [
+                {
+                    "entity": entity["label"],
+                    "word": entity["text"],
+                    "start": entity["start"],
+                    "end": entity["end"],
+                    "score": 0,
+                }
+                for entity in updated_entities
+            ],
+        }
     else:
         model = meal_model
-    return {
-        "text": text,
-        "entities": [
-            {
-                "entity": entity["label"],
-                "word": entity["text"],
-                "start": entity["start"],
-                "end": entity["end"],
-                "score": 0,
-            }
-            for entity in model.predict_entities(
-                text, labels, flat_ner=not nested_ner, threshold=threshold
-            )
-        ],
-    }
+        return {
+            "text": text,
+            "entities": [
+                {
+                    "entity": entity["label"],
+                    "word": entity["text"],
+                    "start": entity["start"],
+                    "end": entity["end"],
+                    "score": 0,
+                }
+                for entity in model.predict_entities(
+                    text, labels, flat_ner=not nested_ner, threshold=threshold
+                )
+            ],
+        }
 
 
 with gr.Blocks(title="GLiNER-M-v2.1") as demo:
@@ -108,7 +154,7 @@ European Championship => competitions
         value=examples[0][0], label="Text input", placeholder="Enter your text here"
     )
     with gr.Row() as row:
-        models = gr.Dropdown((
+        models = gr.Dropdown(
             choices=["ner_meals", "ner_jp"]
             value="ner_jp",
             label="Models",