add model mmBERT

app.py

@@ -101,6 +101,15 @@ custom_css = """
     font-size: 0.9rem;
     margin-bottom: 0;
 }
+.model-header-mmbert {
+    background: linear-gradient(135deg, #f97316, #c2410c);
+    color: white;
+    padding: 8px 16px;
+    border-radius: 8px 8px 0 0;
+    font-weight: bold;
+    font-size: 0.9rem;
+    margin-bottom: 0;
+}
 .batch-sample {
     border: 1px solid #e5e7eb;
     border-radius: 10px;
@@ -168,6 +177,10 @@ custom_css = """
     background: #ede9fe;
     color: #5b21b6;
 }
+.batch-model-label.mmbert {
+    background: #ffedd5;
+    color: #c2410c;
+}
 .dark .batch-model-label.gliner {
     background: #1e3a5f;
     color: #93c5fd;
@@ -176,6 +189,10 @@ custom_css = """
     background: #2e1065;
     color: #c4b5fd;
 }
+.dark .batch-model-label.mmbert {
+    background: #431407;
+    color: #fdba74;
+}
 .batch-entity-tags {
     display: flex;
     flex-wrap: wrap;
@@ -959,11 +976,18 @@ _MODEL_FACTORIES = {
         threshold=0.5,
         use_rules=False,
     ),
+    "mmBERT-multilingual": lambda: mmBERTWrapper(
+        hf_repo="quynong/multilingual",
+        subfolder="jhu-clsp-mmBERT-base",
+        threshold=0.5,
+        use_rules=False,
+    ),
 }
 
 MODEL_DISPLAY_NAMES = {
     "multi_1": "GLiNER2 — multi-v7 (ep3)",
     "XLMBERT": "XLM-RoBERTa (BERT)",
+    "mmBERT-multilingual": "mmBERT-multilingual",
 }
 
 
@@ -1014,12 +1038,13 @@ def run_single_model(text: str, model_key: str):
 
 
 # ==========================================
-# 6. Wrapper cho Gradio (2 model song song)
+# 6. Wrapper cho Gradio (3 model song song)
 # ==========================================
 def process_both(text: str):
     html1, json1 = run_single_model(text, "multi_1")
     html2, json2 = run_single_model(text, "XLMBERT")
-    return html1, json1, html2, json2
+    html3, json3 = run_single_model(text, "mmBERT-multilingual")
+    return html1, json1, html2, json2, html3, json3
 
 
 # ==========================================
@@ -1055,10 +1080,12 @@ def process_test_file(file, model_choice: str):
         return "<div style='color:red;padding:20px;'>File JSON phải là một mảng (array) các object có trường \"text\".</div>", []
 
     model_keys = []
-    if model_choice in ("GLiNER2", "Cả hai"):
+    if model_choice in ("GLiNER2", "Tất cả"):
         model_keys.append("multi_1")
-    if model_choice in ("XLM-RoBERTa", "Cả hai"):
+    if model_choice in ("XLM-RoBERTa", "Tất cả"):
         model_keys.append("XLMBERT")
+    if model_choice in ("mmBERT-multilingual", "Tất cả"):
+        model_keys.append("mmBERT-multilingual")
 
     total = len(data)
     all_results = []
@@ -1067,6 +1094,7 @@ def process_test_file(file, model_choice: str):
     # Summary header placeholder
     total_pii_gliner = 0
     total_pii_bert = 0
+    total_pii_mmbert = 0
 
     for idx, item in enumerate(data):
         text = item.get("text", "")
@@ -1081,7 +1109,7 @@ def process_test_file(file, model_choice: str):
         result_entry = {"sample": sample_num, "text_preview": text[:100]}
 
         # Run models
-        html_gliner = preds_gliner = html_bert = preds_bert = None
+        html_gliner = preds_gliner = html_bert = preds_bert = html_mmbert = preds_mmbert = None
         if "multi_1" in model_keys:
             html_gliner, preds_gliner = run_single_model(text, "multi_1")
             result_entry["gliner"] = preds_gliner
@@ -1092,6 +1120,11 @@ def process_test_file(file, model_choice: str):
             result_entry["xlmbert"] = preds_bert
             if isinstance(preds_bert, list):
                 total_pii_bert += len(preds_bert)
+        if "mmBERT-multilingual" in model_keys:
+            html_mmbert, preds_mmbert = run_single_model(text, "mmBERT-multilingual")
+            result_entry["mmbert_multilingual"] = preds_mmbert
+            if isinstance(preds_mmbert, list):
+                total_pii_mmbert += len(preds_mmbert)
 
         all_results.append(result_entry)
 
@@ -1105,34 +1138,35 @@ def process_test_file(file, model_choice: str):
         )
         html_parts.append(f'<div class="batch-sample-input">{short_text}</div>')
 
-        if len(model_keys) == 2:
-            html_parts.append('<div class="batch-models-row">')
-            # GLiNER column
-            html_parts.append('<div class="batch-model-col">')
-            html_parts.append('<div class="batch-model-label gliner">GLiNER2</div>')
-            if html_gliner:
-                html_parts.append(html_gliner)
-            n_g = len(preds_gliner) if isinstance(preds_gliner, list) else 0
-            html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n_g} entities:</strong>')
-            html_parts.append(_build_entity_tags_html(preds_gliner))
-            html_parts.append('</div></div>')
-            # BERT column
-            html_parts.append('<div class="batch-model-col">')
-            html_parts.append('<div class="batch-model-label bert">XLM-RoBERTa</div>')
-            if html_bert:
-                html_parts.append(html_bert)
-            n_b = len(preds_bert) if isinstance(preds_bert, list) else 0
-            html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n_b} entities:</strong>')
-            html_parts.append(_build_entity_tags_html(preds_bert))
-            html_parts.append('</div></div>')
+        if len(model_keys) >= 2:
+            cols = f"repeat({len(model_keys)}, 1fr)"
+            html_parts.append(f'<div class="batch-models-row" style="grid-template-columns:{cols};">')
+            _batch_cols = [
+                ("multi_1", "gliner", "GLiNER2", html_gliner, preds_gliner),
+                ("XLMBERT", "bert", "XLM-RoBERTa", html_bert, preds_bert),
+                ("mmBERT-multilingual", "mmbert", "mmBERT-multi", html_mmbert, preds_mmbert),
+            ]
+            for mk, cls, name, h, p in _batch_cols:
+                if mk not in model_keys:
+                    continue
+                html_parts.append('<div class="batch-model-col">')
+                html_parts.append(f'<div class="batch-model-label {cls}">{name}</div>')
+                if h:
+                    html_parts.append(h)
+                n = len(p) if isinstance(p, list) else 0
+                html_parts.append(f'<div class="batch-entity-tags"><strong style="font-size:0.75rem;margin-right:6px;">{n} entities:</strong>')
+                html_parts.append(_build_entity_tags_html(p))
+                html_parts.append('</div></div>')
             html_parts.append('</div>')  # end batch-models-row
         else:
             # Single model
             mk = model_keys[0]
-            h = html_gliner if mk == "multi_1" else html_bert
-            p = preds_gliner if mk == "multi_1" else preds_bert
-            cls = "gliner" if mk == "multi_1" else "bert"
-            name = "GLiNER2" if mk == "multi_1" else "XLM-RoBERTa"
+            _single_map = {
+                "multi_1": ("gliner", "GLiNER2", html_gliner, preds_gliner),
+                "XLMBERT": ("bert", "XLM-RoBERTa", html_bert, preds_bert),
+                "mmBERT-multilingual": ("mmbert", "mmBERT-multi", html_mmbert, preds_mmbert),
+            }
+            cls, name, h, p = _single_map[mk]
             html_parts.append(f'<div style="padding:12px 16px;">')
             html_parts.append(f'<div class="batch-model-label {cls}">{name}</div>')
             if h:
@@ -1152,7 +1186,9 @@ def process_test_file(file, model_choice: str):
     if "multi_1" in model_keys:
         summary_lines.append(f"🤖 <strong>GLiNER2:</strong> {total_pii_gliner} PII entities phát hiện<br>")
     if "XLMBERT" in model_keys:
-        summary_lines.append(f"🧠 <strong>XLM-RoBERTa:</strong> {total_pii_bert} PII entities phát hiện")
+        summary_lines.append(f"🧠 <strong>XLM-RoBERTa:</strong> {total_pii_bert} PII entities phát hiện<br>")
+    if "mmBERT-multilingual" in model_keys:
+        summary_lines.append(f"🔶 <strong>mmBERT-multilingual:</strong> {total_pii_mmbert} PII entities phát hiện")
     summary_lines.append("</div>")
 
     final_html = "".join(summary_lines) + "".join(html_parts)
@@ -1175,15 +1211,15 @@ DEFAULT_TEXT = (
     "our support team at 555-0198.\n\n"
 )
 
-with gr.Blocks(title="PII Extractor — So sánh 2 Model", theme=gr.themes.Soft(), css=custom_css) as demo:
-    gr.Markdown("# 🔍 Hệ thống Trích xuất PII — So sánh 2 Model Song Song")
+with gr.Blocks(title="PII Extractor — So sánh 3 Model", theme=gr.themes.Soft(), css=custom_css) as demo:
+    gr.Markdown("# 🔍 Hệ thống Trích xuất PII — So sánh 3 Model Song Song")
 
     with gr.Tabs():
         # ── Tab 1: Nhập văn bản ──────────────────────────────────────
         with gr.TabItem("📝 Nhập văn bản"):
             gr.Markdown(
                 "Nhập văn bản bên dưới và nhấn **Trích xuất PII**. "
-                "Kết quả của **GLiNER2 multi-v7** và **XLM-RoBERTa** sẽ hiển thị đồng thời."
+                "Kết quả của **GLiNER2 multi-v7**, **XLM-RoBERTa** và **mmBERT-multilingual** sẽ hiển thị đồng thời."
             )
 
             with gr.Row():
@@ -1194,7 +1230,7 @@ with gr.Blocks(title="PII Extractor — So sánh 2 Model", theme=gr.themes.Soft(
                     scale=1,
                 )
 
-            submit_btn = gr.Button("⚡ Trích xuất PII (Cả 2 Model)", variant="primary", size="lg")
+            submit_btn = gr.Button("⚡ Trích xuất PII (Cả 3 Model)", variant="primary", size="lg")
 
             gr.Markdown("## 📊 Kết quả nhận diện")
             with gr.Row(equal_height=True):
@@ -1210,6 +1246,12 @@ with gr.Blocks(title="PII Extractor — So sánh 2 Model", theme=gr.themes.Soft(
                     )
                     out_html_bert = gr.HTML(label="XLM-RoBERTa Output")
 
+                with gr.Column():
+                    gr.HTML(
+                        '<div class="model-header-mmbert">🔶 Model 3 — mmBERT-multilingual (jhu-clsp, threshold=0.5)</div>'
+                    )
+                    out_html_mmbert = gr.HTML(label="mmBERT-multilingual Output")
+
             gr.Markdown("## 🗂️ Dữ liệu JSON thô")
             with gr.Row(equal_height=True):
                 with gr.Column():
@@ -1220,10 +1262,14 @@ with gr.Blocks(title="PII Extractor — So sánh 2 Model", theme=gr.themes.Soft(
                     gr.Markdown("**XLM-RoBERTa JSON**")
                     out_json_bert = gr.JSON(label="XLM-RoBERTa Raw JSON")
 
+                with gr.Column():
+                    gr.Markdown("**mmBERT-multilingual JSON**")
+                    out_json_mmbert = gr.JSON(label="mmBERT-multilingual Raw JSON")
+
             submit_btn.click(
                 fn=process_both,
                 inputs=[input_text],
-                outputs=[out_html_gliner, out_json_gliner, out_html_bert, out_json_bert],
+                outputs=[out_html_gliner, out_json_gliner, out_html_bert, out_json_bert, out_html_mmbert, out_json_mmbert],
             )
 
         # ── Tab 2: Test File ─────────────────────────────────────────
@@ -1242,8 +1288,8 @@ with gr.Blocks(title="PII Extractor — So sánh 2 Model", theme=gr.themes.Soft(
                     )
                 with gr.Column(scale=1):
                     model_choice = gr.Radio(
-                        choices=["Cả hai", "GLiNER2", "XLM-RoBERTa"],
-                        value="Cả hai",
+                        choices=["Tất cả", "GLiNER2", "XLM-RoBERTa", "mmBERT-multilingual"],
+                        value="Tất cả",
                         label="🔧 Chọn model",
                     )