Update app.py

app.py

@@ -22,9 +22,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 model.eval()
 
-def merge_bio_spans(text: str, per_token):
+def merge_spans(text: str, per_token):
     """
     per_token: list of dict {label, start, end}
+    Supports:
+      - BIO labels: B-XXX / I-XXX / O
+      - Non-BIO labels: XXX / O
     returns: list of dict {entity, start, end}
     """
     spans = []
@@ -40,13 +43,37 @@ def merge_bio_spans(text: str, per_token):
         lab = t["label"]
         st, ed = t["start"], t["end"]
 
+        # normalize
+        if lab is None:
+            lab = "O"
+
+        if lab == "O":
+            close_cur()
+            continue
+
+        # BIO case
         if lab.startswith("B-"):
             close_cur()
             cur = {"entity": lab[2:], "start": st, "end": ed}
-        elif lab.startswith("I-") and cur and cur["entity"] == lab[2:]:
+            continue
+
+        if lab.startswith("I-"):
+            ent = lab[2:]
+            if cur and cur["entity"] == ent:
+                cur["end"] = ed
+            else:
+                # treat as a new span if I- appears without proper B-
+                close_cur()
+                cur = {"entity": ent, "start": st, "end": ed}
+            continue
+
+        # Non-BIO case: label like "person" / "ORG" / etc.
+        ent = lab
+        if cur and cur["entity"] == ent:
             cur["end"] = ed
         else:
             close_cur()
+            cur = {"entity": ent, "start": st, "end": ed}
 
     close_cur()
     return spans
@@ -81,8 +108,8 @@ def run_ner(text: str, max_length: int, show_tokens: bool):
     # build per-token labels (skip specials)
     per_token = []
     for tok, pid, (st, ed) in zip(tokens, pred_ids, offsets):
-        if tok in tokenizer.all_special_tokens:
-            continue
+        ＃if tok in tokenizer.all_special_tokens:
+        ＃    continue
         if st == ed:
             continue
         per_token.append({
@@ -92,7 +119,7 @@ def run_ner(text: str, max_length: int, show_tokens: bool):
             "end": int(ed),
         })
 
-    spans = merge_bio_spans(text, per_token)
+    spans = merge_spans(text, per_token)
 
     # Return 2D list to avoid `[object Object]`
     table_rows = []