fix highlighting

app.py

@@ -112,7 +112,7 @@ def prune_acronym_and_self_relations(ner_preds, rel_preds):
 # Highlighting function
 
 def highlight_text(text, ner_threshold, rel_threshold):
-    # Run inference
+    # 1) Inference
     ner_preds, rel_preds = inference_pipeline(
         text,
         model=model,
@@ -124,32 +124,43 @@ def highlight_text(text, ner_threshold, rel_threshold):
         re_multi_label=False,
         return_index=True,
     )
-
-    # Post-process
     ner_preds, rel_preds = prune_acronym_and_self_relations(ner_preds, rel_preds)
 
-    # Gather all spans
+    # 2) Compute how long the RE prompt prefix is
+    #    This must match exactly what your extractor prepends:
+    prefix = f"{relation_extractor.prompt} \n "
+    prefix_len = len(prefix)
+
+    # 3) Gather spans
     spans = []
     for ent in ner_preds:
         spans.append((ent["start"], ent["end"], ent["label"]))
+
+    #    Use the extractor‐returned start/end, minus prefix_len
     for src, rels in rel_preds.items():
         for r in rels:
-            for m in re.finditer(re.escape(r["target"]), text):
-                spans.append((m.start(), m.end(), f"{src} <> {r['relation']}"))
-
-    # Merge labels by span
+            # adjust the indices back onto the raw text
+            s = r["start"] - prefix_len
+            e = r["end"]   - prefix_len
+            # skip anything that fell outside
+            if s < 0 or e < 0:
+                continue
+            label = f"{r['source']} <> {r['relation']}"
+            spans.append((s, e, label))
+
+    # 4) Merge & build entities (same as before)
     merged = defaultdict(list)
-    for start, end, lbl in spans:
-        merged[(start, end)].append(lbl)
+    for s, e, lbl in spans:
+        merged[(s, e)].append(lbl)
 
-    # Build Gradio entities
     entities = []
-    for (start, end), lbls in sorted(merged.items(), key=lambda x: x[0]):
+    for (s, e), lbls in sorted(merged.items(), key=lambda x: x[0]):
         entities.append({
             "entity": ", ".join(lbls),
-            "start": start,
-            "end": end
+            "start":  s,
+            "end":    e
         })
+
     return {"text": text, "entities": entities}, {"ner": ner_preds, "relations": rel_preds}
 
 # JSON output function