Update app.py

app.py

@@ -7,20 +7,36 @@ from key_store import (
 )
 
 # ── Load your private HF model ────────────────────────────────────────────────
-HF_TOKEN = os.environ.get("HF_TOKEN")  # set as a Space secret
+HF_TOKEN = os.environ.get("HF_TOKEN")
 MODEL_ID  = "Wardline/CLOAK_V2.2.7"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 model     = AutoModelForTokenClassification.from_pretrained(MODEL_ID, token=HF_TOKEN)
 ner_pipe  = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
+# ── Stitcher: merge split sub-word tokens ─────────────────────────────────────
+def stitch_entities(raw_entities: list) -> list:
+    merged = []
+    for r in raw_entities:
+        if (merged
+                and merged[-1]["entity_group"] == r["entity_group"]
+                and merged[-1]["end"] == r["start"]):
+            merged[-1]["word"] += r["word"].replace(" ", "")
+            merged[-1]["end"]   = r["end"]
+        else:
+            merged.append(dict(r))  # copy so we don't mutate the original
+    return merged
+
 # ── PII detection ─────────────────────────────────────────────────────────────
 def detect_pii(text: str, api_key: str) -> dict:
     entry = validate_key(api_key)
     if not entry:
         return {"error": "Invalid or revoked API key.", "entities": []}
     increment_usage(api_key)
-    entities = ner_pipe(text)
+
+    raw      = ner_pipe(text)
+    entities = stitch_entities(raw)
+
     return {
         "text": text,
         "entities": [
@@ -29,10 +45,10 @@ def detect_pii(text: str, api_key: str) -> dict:
                 "word":   e["word"],
                 "score":  round(e["score"], 4),
                 "start":  e["start"],
-                "end":    e["end"]
+                "end":    e["end"],
             }
             for e in entities
-        ]
+        ],
     }
 
 # ── Key management helpers ────────────────────────────────────────────────────
@@ -50,11 +66,11 @@ def get_key_table():
 
 def do_revoke(full_key):
     revoke_key(full_key.strip())
-    return f"Key revoked.", get_key_table()
+    return "Key revoked.", get_key_table()
 
 def do_delete(full_key):
     delete_key(full_key.strip())
-    return f"Key deleted.", get_key_table()
+    return "Key deleted.", get_key_table()
 
 # ── Gradio UI ─────────────────────────────────────────────────────────────────
 with gr.Blocks(title="PII Detection API") as demo:
@@ -65,18 +81,14 @@ with gr.Blocks(title="PII Detection API") as demo:
             api_key_in = gr.Textbox(label="API key", placeholder="pii_sk_...")
             text_in    = gr.Textbox(label="Input text", lines=4,
                                     placeholder="Enter text to scan for PII...")
-        detect_btn = gr.Button("Detect PII", variant="primary")
+        detect_btn  = gr.Button("Detect PII", variant="primary")
         output_json = gr.JSON(label="Results")
-        detect_btn.click(
-            fn=detect_pii,
-            inputs=[text_in, api_key_in],
-            outputs=output_json
-        )
+        detect_btn.click(fn=detect_pii, inputs=[text_in, api_key_in], outputs=output_json)
 
     with gr.Tab("API Key Management"):
         gr.Markdown("### Create a new key")
         with gr.Row():
-            label_in  = gr.Textbox(label="Key label", placeholder="e.g. production")
+            label_in   = gr.Textbox(label="Key label", placeholder="e.g. production")
             create_btn = gr.Button("Generate key")
         create_msg = gr.Textbox(label="New key (copy now!)", interactive=False, lines=2)
 
@@ -84,7 +96,7 @@ with gr.Blocks(title="PII Detection API") as demo:
         key_df = gr.DataFrame(
             value=get_key_table(),
             headers=["Label", "Key preview", "Status", "Usage", "Created"],
-            interactive=False
+            interactive=False,
         )
 
         gr.Markdown("### Revoke or delete a key")
@@ -94,8 +106,8 @@ with gr.Blocks(title="PII Detection API") as demo:
             delete_btn    = gr.Button("Delete", variant="stop")
         manage_msg = gr.Textbox(label="Result", interactive=False)
 
-        create_btn.click(create_key, inputs=label_in,       outputs=[create_msg, key_df])
-        revoke_btn.click(do_revoke,  inputs=manage_key_in,  outputs=[manage_msg, key_df])
-        delete_btn.click(do_delete,  inputs=manage_key_in,  outputs=[manage_msg, key_df])
+        create_btn.click(create_key, inputs=label_in,      outputs=[create_msg, key_df])
+        revoke_btn.click(do_revoke,  inputs=manage_key_in, outputs=[manage_msg, key_df])
+        delete_btn.click(do_delete,  inputs=manage_key_in, outputs=[manage_msg, key_df])
 
 demo.launch()