Update app.py

app.py

@@ -42,6 +42,286 @@ EXAMPLES = [
     },
 ]
 
+# ── Custom CSS ──────────────────────────────────────────────────────────────
+
+CUSTOM_CSS = """
+/* ── Global ──────────────────────────────────────── */
+.gradio-container {
+    max-width: 1200px !important;
+    font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif !important;
+}
+
+/* ── Header Banner ───────────────────────────────── */
+.header-banner {
+    background: linear-gradient(135deg, #1a1a2e 0%, #16213e 40%, #0f3460 100%);
+    border-radius: 16px;
+    padding: 32px 36px;
+    margin-bottom: 24px;
+    color: #fff;
+    position: relative;
+    overflow: hidden;
+}
+.header-banner::before {
+    content: '';
+    position: absolute;
+    top: -50%;
+    right: -20%;
+    width: 400px;
+    height: 400px;
+    background: radial-gradient(circle, rgba(99, 179, 237, 0.15) 0%, transparent 70%);
+    pointer-events: none;
+}
+.header-banner h1 {
+    margin: 0 0 8px 0 !important;
+    font-size: 28px !important;
+    font-weight: 700 !important;
+    letter-spacing: -0.5px;
+    background: linear-gradient(90deg, #ffffff, #a8d8ea);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+}
+.header-banner p {
+    margin: 0 !important;
+    font-size: 14px !important;
+    color: rgba(255, 255, 255, 0.7) !important;
+    line-height: 1.6;
+}
+.header-badge {
+    display: inline-block;
+    background: rgba(99, 179, 237, 0.2);
+    color: #a8d8ea;
+    font-size: 11px;
+    font-weight: 600;
+    padding: 3px 10px;
+    border-radius: 20px;
+    margin-bottom: 12px;
+    border: 1px solid rgba(99, 179, 237, 0.3);
+    letter-spacing: 0.5px;
+    text-transform: uppercase;
+}
+
+/* ── Section Cards ───────────────────────────────── */
+.section-card {
+    background: #ffffff;
+    border: 1px solid #e8ecf1;
+    border-radius: 12px;
+    padding: 20px;
+    margin-bottom: 16px;
+    box-shadow: 0 1px 3px rgba(0, 0, 0, 0.04);
+    transition: box-shadow 0.2s ease;
+}
+.section-card:hover {
+    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.06);
+}
+.section-title {
+    font-size: 13px !important;
+    font-weight: 600 !important;
+    color: #64748b !important;
+    text-transform: uppercase !important;
+    letter-spacing: 0.8px !important;
+    margin-bottom: 12px !important;
+    padding-bottom: 8px !important;
+    border-bottom: 1px solid #f1f5f9 !important;
+}
+
+/* ── Form Controls ───────────────────────────────── */
+.gr-textbox, .gr-dropdown {
+    border-radius: 8px !important;
+    border: 1px solid #e2e8f0 !important;
+    transition: border-color 0.2s ease, box-shadow 0.2s ease !important;
+}
+.gr-textbox:focus-within, .gr-dropdown:focus-within {
+    border-color: #63b3ed !important;
+    box-shadow: 0 0 0 3px rgba(99, 179, 237, 0.12) !important;
+}
+textarea {
+    font-size: 13.5px !important;
+    line-height: 1.6 !important;
+}
+
+/* ── Document List ───────────────────────────────── */
+.doc-list-container .gr-textbox {
+    border-left: 3px solid #63b3ed !important;
+    margin-bottom: 8px !important;
+}
+
+/* ── Buttons ─────────────────────────────────────── */
+.doc-btn {
+    border-radius: 8px !important;
+    font-weight: 500 !important;
+    font-size: 13px !important;
+    padding: 8px 16px !important;
+    transition: all 0.2s ease !important;
+}
+.btn-add {
+    background: #f0f7ff !important;
+    color: #2b6cb0 !important;
+    border: 1px solid #bee3f8 !important;
+}
+.btn-add:hover {
+    background: #dbeafe !important;
+    border-color: #90cdf4 !important;
+}
+.btn-remove {
+    background: #fff5f5 !important;
+    color: #c53030 !important;
+    border: 1px solid #fed7d7 !important;
+}
+.btn-remove:hover {
+    background: #fee2e2 !important;
+    border-color: #feb2b2 !important;
+}
+
+/* ── Submit Button Override ──────────────────────── */
+.submit-row .primary-btn {
+    width: 100%;
+    padding: 12px !important;
+    font-size: 15px !important;
+    font-weight: 600 !important;
+    border-radius: 10px !important;
+    background: linear-gradient(135deg, #2b6cb0, #3182ce) !important;
+    border: none !important;
+    letter-spacing: 0.3px;
+    transition: all 0.25s ease !important;
+}
+.submit-row .primary-btn:hover {
+    background: linear-gradient(135deg, #2c5282, #2b6cb0) !important;
+    box-shadow: 0 4px 14px rgba(43, 108, 176, 0.35) !important;
+    transform: translateY(-1px);
+}
+
+/* ── Results Panel ──────────���────────────────────── */
+.result-panel {
+    background: #fafbfc;
+    border: 1px solid #e8ecf1;
+    border-radius: 12px;
+    padding: 20px;
+    min-height: 200px;
+}
+.result-panel table {
+    border-collapse: separate !important;
+    border-spacing: 0 !important;
+    width: 100% !important;
+    font-size: 13px !important;
+}
+.result-panel thead th {
+    background: #f1f5f9 !important;
+    color: #475569 !important;
+    font-weight: 600 !important;
+    font-size: 11px !important;
+    text-transform: uppercase !important;
+    letter-spacing: 0.5px !important;
+    padding: 10px 12px !important;
+    border-bottom: 2px solid #e2e8f0 !important;
+}
+.result-panel thead th:first-child {
+    border-radius: 8px 0 0 0 !important;
+}
+.result-panel thead th:last-child {
+    border-radius: 0 8px 0 0 !important;
+}
+.result-panel tbody td {
+    padding: 10px 12px !important;
+    border-bottom: 1px solid #f1f5f9 !important;
+    vertical-align: top !important;
+    color: #334155 !important;
+}
+.result-panel tbody tr:hover {
+    background: #f8fafc !important;
+}
+.result-panel tbody tr:last-child td {
+    border-bottom: none !important;
+}
+
+/* ── Summary Badge ───────────────────────────────── */
+.summary-badge {
+    background: linear-gradient(135deg, #f0fff4, #e6fffa);
+    border: 1px solid #c6f6d5;
+    border-radius: 10px;
+    padding: 12px 16px;
+    margin-top: 12px;
+    font-size: 13px;
+    color: #276749;
+}
+
+/* ── Citation Accordion ──────────────────────────── */
+.citation-box {
+    background: #f8fafc;
+    border: 1px solid #e2e8f0;
+    border-radius: 10px;
+    padding: 16px 20px;
+    margin-top: 20px;
+}
+.citation-box pre {
+    background: #1e293b !important;
+    border-radius: 8px !important;
+    padding: 14px !important;
+    font-size: 11.5px !important;
+    line-height: 1.6 !important;
+    color: #e2e8f0 !important;
+    overflow-x: auto;
+}
+.citation-box h3 {
+    font-size: 14px !important;
+    color: #64748b !important;
+    margin: 0 0 10px 0 !important;
+}
+
+/* ── Config Row (Instruction + ChunkSize + Example) */
+.config-row {
+    display: flex;
+    gap: 12px;
+    align-items: flex-end;
+}
+.config-row > div {
+    flex: 1;
+}
+
+/* ── Scrollbar ───────────────────────────────────── */
+.doc-scroll::-webkit-scrollbar {
+    width: 5px;
+}
+.doc-scroll::-webkit-scrollbar-track {
+    background: transparent;
+}
+.doc-scroll::-webkit-scrollbar-thumb {
+    background: #cbd5e1;
+    border-radius: 10px;
+}
+.doc-scroll::-webkit-scrollbar-thumb:hover {
+    background: #94a3b8;
+}
+
+/* ── Label tweaks ────────────────────────────────── */
+.label-query label span {
+    font-size: 14px !important;
+    font-weight: 600 !important;
+    color: #1e293b !important;
+}
+"""
+
+# ── Theme ───────────────────────────────────────────────────────────────────
+
+theme = gr.themes.Soft(
+    primary_hue="blue",
+    secondary_hue="slate",
+    neutral_hue="slate",
+    font=gr.themes.GoogleFont("Inter"),
+).set(
+    body_background_fill="#f8fafc",
+    body_background_fill_dark="#0f172a",
+    background_fill_primary="#ffffff",
+    background_fill_secondary="#f1f5f9",
+    border_color_primary="#e2e8f0",
+    text_primary="#1e293b",
+    text_secondary="#64748b",
+    radius_quarter="8px",
+    radius_half="10px",
+)
+
+# ── Logic helpers ───────────────────────────────────────────────────────────
+
 
 def _make_vis_updates(count: int):
     return [gr.update(visible=(i < count)) for i in range(MAX_DOCS)]
@@ -96,7 +376,7 @@ def rerank(query: str, instruction: str, chunk_size: int, *doc_texts: str):
         chunk_size=chunk_size,
     )
 
-#    query = query.strip()
+    query = query.strip()
     instruction = instruction.strip() or DEFAULT_INSTRUCTION
 
     try:
@@ -114,8 +394,9 @@ def rerank(query: str, instruction: str, chunk_size: int, *doc_texts: str):
             table.append([rank_idx, corpus_id, round(score, 6), doc])
 
         summary = (
-            f"Reranked {len(docs)} documents with "
-            f"`{MODEL_ID}` (chunk_size={chunk_size}). Higher score means more relevant."
+            f"Reranked **{len(docs)}** documents with "
+            f"`{MODEL_ID}` (chunk_size={chunk_size}).  "
+            f"Higher score means more relevant."
         )
         return table, summary
     except Exception as error:
@@ -124,72 +405,138 @@ def rerank(query: str, instruction: str, chunk_size: int, *doc_texts: str):
 
 # ── UI ──────────────────���───────────────────────────────────────────────────
 
-with gr.Blocks(title="KaLM-Reranker-V1 Demo") as demo:
-    gr.Markdown(
-        """
-# KaLM-Reranker-V1 Demo
-**KaLM-Reranker-V1** is a fast but not late-interaction reranker for compressed document reranking.
-Input a query and several candidate documents. The demo returns relevance scores and reranked results.
-        """
-    )
-
-    with gr.Row():
+with gr.Blocks(
+    title="KaLM-Reranker-V1 Demo",
+    theme=theme,
+    css=CUSTOM_CSS,
+) as demo:
+
+    # ── Header ──
+    gr.HTML("""
+    <div class="header-banner">
+        <div class="header-badge">Powered by KaLM-Embedding</div>
+        <h1>KaLM-Reranker-V1 Demo</h1>
+        <p>A fast but not late-interaction reranker for compressed document reranking. <br/>
+        Input a query and candidate documents to get relevance scores and reranked results.</p>
+    </div>
+    """)
+
+    with gr.Row(equal_height=False):
         # ── Left column: inputs ──
         with gr.Column(scale=1):
-            query_input = gr.Textbox(
-                label="Query",
-                value="",
-                lines=2,
-            )
-            instruction_input = gr.Textbox(
-                label="Instruction",
-                value=DEFAULT_INSTRUCTION,
-                lines=2,
-            )
-            chunk_size_input = gr.Dropdown(
-                label="Chunk Size",
-                choices=[2, 4, 8, 16],
-                value=2,
-                interactive=True,
-            )
-
-            # Example selector
-            example_selector = gr.Dropdown(
-                label="Load Example",
-                choices=[ex["label"] for ex in EXAMPLES],
-                value=None,
-                interactive=True,
-            )
-
-            gr.Markdown("### Candidate Documents")
-            doc_count = gr.State(value=INIT_DOCS)
-            doc_textboxes: List[gr.Textbox] = []
-
-            with gr.Column() as doc_list:
-                for i in range(MAX_DOCS):
-                    visible = i < INIT_DOCS
-                    tb = gr.Textbox(
-                        label=f"Document {i + 1}",
-                        lines=2,
-                        placeholder=f"Enter document {i + 1}...",
-                        visible=visible,
-                    )
-                    doc_textboxes.append(tb)
 
-            with gr.Row():
-                add_btn = gr.Button("+ Add Document", size="sm")
-                remove_btn = gr.Button("- Remove Last", size="sm")
+            # Query
+            with gr.Group(elem_classes=["section-card"]):
+                gr.Markdown('<div class="section-title">Query</div>', elem_classes="label-query")
+                query_input = gr.Textbox(
+                    show_label=False,
+                    value="What is the capital of China?",
+                    lines=3,
+                    placeholder="Enter your query here...",
+                )
+
+            # Config row: Instruction · Chunk Size · Load Example
+            with gr.Group(elem_classes=["section-card"]):
+                gr.Markdown('<div class="section-title">Configuration</div>', elem_classes="label-query")
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        instruction_input = gr.Textbox(
+                            label="Instruction",
+                            value=DEFAULT_INSTRUCTION,
+                            lines=2,
+                            show_copy_button=True,
+                        )
+                    with gr.Column(scale=1):
+                        chunk_size_input = gr.Dropdown(
+                            label="Chunk Size",
+                            choices=[2, 4, 8, 16],
+                            value=4,
+                            interactive=True,
+                        )
+                example_selector = gr.Dropdown(
+                    label="Quick Example",
+                    choices=[""] + [ex["label"] for ex in EXAMPLES],
+                    value="",
+                    interactive=True,
+                    allow_custom_value=False,
+                )
+
+            # Candidate Documents
+            with gr.Group(elem_classes=["section-card"]):
+                gr.Markdown(
+                    '<div class="section-title">Candidate Documents '
+                    '<span style="float:right;font-weight:400;text-transform:none;letter-spacing:0;font-size:11px;color:#94a3b8">'
+                    f'Max {MAX_DOCS} docs</span></div>',
+                    elem_classes="label-query",
+                )
+                doc_count = gr.State(value=INIT_DOCS)
+                doc_textboxes: List[gr.Textbox] = []
+
+                with gr.Column(elem_classes=["doc-list-container", "doc-scroll"]) as doc_list:
+                    for i in range(MAX_DOCS):
+                        visible = i < INIT_DOCS
+                        tb = gr.Textbox(
+                            label=f"  {i + 1}",
+                            lines=2,
+                            placeholder=f"Enter document {i + 1}...",
+                            visible=visible,
+                            show_label=True,
+                        )
+                        doc_textboxes.append(tb)
+
+                with gr.Row():
+                    add_btn = gr.Button(
+                        "+  Add Document",
+                        size="sm",
+                        elem_classes=["doc-btn", "btn-add"],
+                    )
+                    remove_btn = gr.Button(
+                        "-  Remove Last",
+                        size="sm",
+                        elem_classes=["doc-btn", "btn-remove"],
+                    )
 
-            submit = gr.Button("Rerank", variant="primary")
+            # Submit
+            with gr.Row(elem_classes=["submit-row"]):
+                submit = gr.Button(
+                    "Rerank",
+                    variant="primary",
+                    elem_classes=["primary-btn"],
+                )
 
         # ── Right column: outputs ──
         with gr.Column(scale=1):
-            output_table = gr.Dataframe(
-                headers=["Rank", "Corpus ID", "Score", "Document"],
-                label="Reranking Results",
-                wrap=True,
-            )
-            output_summary = gr.Markdown()
+            with gr.Group(elem_classes=["result-panel"]):
+                gr.Markdown(
+                    '<div class="section-title">Reranking Results</div>',
+                    elem_classes="label-query",
+                )
+                output_table = gr.Dataframe(
+                    headers=["Rank", "Corpus ID", "Score", "Document"],
+                    label=None,
+                    wrap=True,
+                    interactive=False,
+                )
+                output_summary = gr.HTML()
+
+    # ── Citation ──
+    gr.HTML("""
+    <details class="citation-box">
+        <summary style="cursor:pointer;font-weight:600;font-size:14px;color:#64748b;">
+            Cite this work
+        </summary>
+        <pre><code>@misc{zhao2026kalmrerankerv1,
+      title={KaLM-Reranker-V1: Fast but Not Late Interaction for Compressed Document Reranking},
+      author={Xinping Zhao and Jiaxin Xu and Ziqi Dai and Xin Zhang and Shouzheng Huang
+              and Danyu Tang and Xinshuo Hu and Meishan Zhang and Baotian Hu and Min Zhang},
+      year={2026},
+      eprint={2606.22807},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2606.22807},
+}</code></pre>
+    </details>
+    """)
 
     # ── Events ──
     add_btn.click(
@@ -213,26 +560,6 @@ Input a query and several candidate documents. The demo returns relevance scores
         outputs=[output_table, output_summary],
     )
 
-    gr.Markdown(
-        """
-## Citation
-
-If you find this demo useful, please cite:
-
-```bibtex
-@misc{zhao2026kalmrerankerv1,
-      title={KaLM-Reranker-V1: Fast but Not Late Interaction for Compressed Document Reranking},
-      author={Xinping Zhao and Jiaxin Xu and Ziqi Dai and Xin Zhang and Shouzheng Huang and Danyu Tang and Xinshuo Hu and Meishan Zhang and Baotian Hu and Min Zhang},
-      year={2026},
-      eprint={2606.22807},
-      archivePrefix={arXiv},
-      primaryClass={cs.CL},
-      url={https://arxiv.org/abs/2606.22807},
-}
-```
-        """
-    )
-
 
 if __name__ == "__main__":
     demo.launch()