NER tab: render entities with gr.HighlightedText

app.py

@@ -40,13 +40,82 @@ def _parse_labels(label_text: str):
     return list(labels.keys())
 
 
+def _to_highlighted(text: str, result: dict, include_confidence: bool = True):
+    """Convert a GLiNER2 entity result into ``gr.HighlightedText`` spans.
+
+    Returns a list of ``(substring, label_or_None)`` tuples covering the whole
+    input text, so unmatched text is rendered plain and matches are highlighted
+    with their entity type.
+    """
+    entities = (result or {}).get("entities") or {}
+    found = []  # (start, end, label, confidence)
+
+    for label, items in entities.items():
+        if not isinstance(items, list):
+            items = [items]
+        cursor = 0
+        for item in items:
+            if isinstance(item, dict):
+                surface = item.get("text", "")
+                start = item.get("start", item.get("char_start"))
+                end = item.get("end", item.get("char_end"))
+                conf = item.get("confidence")
+            else:
+                surface, start, end, conf = str(item), None, None, None
+            if start is not None and end is not None:
+                start, end = int(start), int(end)
+                # Trim whitespace that token boundaries may have swallowed.
+                while start < end and text[start].isspace():
+                    start += 1
+                while end > start and text[end - 1].isspace():
+                    end -= 1
+                if surface and text[start:end].strip() != surface.strip():
+                    start = end = None  # offsets don't match: fall back to search
+            if start is None or end is None:
+                # Fall back to locating the surface form in the text.
+                if not surface:
+                    continue
+                start = text.find(surface, cursor)
+                if start == -1:
+                    start = text.find(surface)
+                    if start == -1:
+                        continue
+                end = start + len(surface)
+                cursor = end
+            if end <= start:
+                continue
+            found.append((int(start), int(end), str(label), conf))
+
+    # Drop overlaps (keep the higher-confidence / longer span) so the text can
+    # be sliced into a clean sequence of segments.
+    found.sort(key=lambda s: (-(s[3] if s[3] is not None else 0.0), -(s[1] - s[0])))
+    kept = []
+    for span in found:
+        if all(span[1] <= k[0] or span[0] >= k[1] for k in kept):
+            kept.append(span)
+    kept.sort(key=lambda s: s[0])
+
+    segments = []
+    index = 0
+    for start, end, label, conf in kept:
+        if start > index:
+            segments.append((text[index:start], None))
+        tag = f"{label} ({conf:.2f})" if include_confidence and conf is not None else label
+        segments.append((text[start:end], tag))
+        index = end
+    if index < len(text):
+        segments.append((text[index:], None))
+    if not segments:
+        segments = [(text, None)]
+    return segments
+
+
 @spaces.GPU(duration=15)
 def extract_entities(
     text: str,
     labels: str,
     threshold: float = 0.5,
     include_confidence: bool = True,
-    include_spans: bool = True,
 ):
     """Extract named entities from text using a user-defined label schema.
 
@@ -55,13 +124,16 @@ def extract_entities(
         labels: Entity types, one per line. Use ``label::description`` for
             richer prompts (e.g. ``person::individual human``).
         threshold: Confidence threshold (0–1). Lower finds more, noisier matches.
-        include_confidence: Include confidence scores in the output.
-        include_spans: Include character offsets in the output.
+        include_confidence: Show confidence scores next to the entity labels.
+
+    Returns:
+        A list of ``(substring, entity_type_or_None)`` tuples for
+        ``gr.HighlightedText``.
     """
     if not text.strip():
-        return json.dumps({"error": "Please enter some text."}, indent=2)
+        return [("Please enter some text.", "error")]
     if not labels.strip():
-        return json.dumps({"error": "Please specify at least one entity type."}, indent=2)
+        return [("Please specify at least one entity type.", "error")]
 
     try:
         entity_types = _parse_labels(labels)
@@ -69,12 +141,12 @@ def extract_entities(
             text,
             entity_types,
             threshold=threshold,
-            include_confidence=include_confidence,
-            include_spans=include_spans,
+            include_confidence=True,
+            include_spans=True,
         )
-        return json.dumps(result, indent=2, default=str)
+        return _to_highlighted(text, result, include_confidence=include_confidence)
     except Exception as e:
-        return json.dumps({"error": str(e)}, indent=2)
+        return [(f"Error: {e}", "error")]
 
 
 @spaces.GPU(duration=15)
@@ -303,11 +375,14 @@ with gr.Blocks(
                             0.0, 1.0, value=0.5, step=0.05,
                             label="Confidence Threshold",
                         )
-                        ner_conf = gr.Checkbox(value=True, label="Include confidence")
-                        ner_spans = gr.Checkbox(value=True, label="Include spans")
+                        ner_conf = gr.Checkbox(value=True, label="Show confidence in labels")
                     ner_btn = gr.Button("Extract Entities", variant="primary")
                 with gr.Column(scale=2):
-                    ner_out = gr.Code(label="Results (JSON)", language="json", lines=18)
+                    ner_out = gr.HighlightedText(
+                        label="Extracted Entities",
+                        combine_adjacent=True,
+                        show_legend=True,
+                    )
             gr.Examples(
                 examples=NER_EXAMPLES,
                 inputs=[ner_text, ner_labels, ner_threshold],
@@ -318,7 +393,7 @@ with gr.Blocks(
             )
             ner_btn.click(
                 fn=extract_entities,
-                inputs=[ner_text, ner_labels, ner_threshold, ner_conf, ner_spans],
+                inputs=[ner_text, ner_labels, ner_threshold, ner_conf],
                 outputs=ner_out,
                 api_name="/extract_entities",
             )