app.py

import spaces
import os
from typing import List
import gradio as gr

os.environ["TOKENIZERS_PARALLELISM"] = "false"
from kalm_reranker import KaLMReranker

MODEL_ID = "KaLM-Embedding/KaLM-Reranker-V1-Nano"
DEFAULT_INSTRUCTION = "Given a query, retrieve documents that answer the query."
MAX_DOCS = 20
MAX_DOC_CHARS = 4000
INIT_DOCS = 4

EXAMPLES = [
    {
        "label": "Capital of China",
        "query": "What is the capital of China?",
        "docs": [
            "The capital of China is Beijing.",
            "Gravity attracts bodies toward one another.",
            "Paris is the capital of France.",
        ],
    },
    {
        "label": "Efficient Reranking",
        "query": "Which model is suitable for efficient reranking?",
        "docs": [
            "KaLM-Reranker-V1-Nano is designed for efficient reranking.",
            "Large language models are often expensive for reranking.",
            "Image classifiers are used for visual recognition.",
        ],
    },
    {
        "label": "KaLM-Reranker Design",
        "query": "What is KaLM-Reranker-V1 designed for?",
        "docs": [
            "KaLM-Reranker-V1 is a reranker for compressed document reranking.",
            "KaLM-Embedding is a general-purpose embedding model.",
            "Weather forecasting predicts future weather conditions.",
        ],
    },
]


def _make_vis_updates(count: int):
    return [gr.update(visible=(i < count)) for i in range(MAX_DOCS)]


def add_doc(count: int):
    count = min(count + 1, MAX_DOCS)
    return _make_vis_updates(count) + [count]


def remove_doc(count: int):
    count = max(count - 1, 1)
    return _make_vis_updates(count) + [count]


def load_example(example_label):
    if not example_label:
        query_upd = gr.update(value="")
        doc_upds = [
            gr.update(value="", visible=(i < INIT_DOCS))
            for i in range(MAX_DOCS)
        ]
        return [query_upd] + doc_upds + [INIT_DOCS]

    ex = next((e for e in EXAMPLES if e["label"] == example_label), None)
    if ex is None:
        return [gr.update()] + [gr.update() for _ in range(MAX_DOCS)] + [INIT_DOCS]

    docs = ex["docs"]
    count = len(docs)
    query_upd = gr.update(value=ex["query"])
    doc_upds = [
        gr.update(value=docs[i] if i < count else "", visible=(i < count))
        for i in range(MAX_DOCS)
    ]
    return [query_upd] + doc_upds + [count]


@spaces.GPU
def rerank(query: str, instruction: str, chunk_size: int, *doc_texts: str):
    docs = [doc.strip() for doc in doc_texts if doc and doc.strip()]
    docs = docs[:MAX_DOCS]
    docs = [doc[:MAX_DOC_CHARS] for doc in docs]

    if not query.strip():
        return [], "Please input a query."
    if not docs:
        return [], "Please input at least one candidate document."

    reranker = KaLMReranker(
        MODEL_ID,
        device=None,
        dtype=None,
        batch_size=4,
        query_max_length=512,
        max_length=1024,
        chunk_size=chunk_size,
    )

    query = query.strip()
    instruction = instruction.strip() or DEFAULT_INSTRUCTION

    try:
        rankings = reranker.rank(
            query=query,
            documents=docs,
            instruction=instruction,
        )

        table = []
        for rank_idx, item in enumerate(rankings, start=1):
            corpus_id = item["corpus_id"]
            score = float(item["score"])
            doc = docs[corpus_id]
            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."
        )
        return table, summary
    except Exception as error:
        return [], f"Error during reranking: {repr(error)}"


# ── 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():
        # ── 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 = gr.Dropdown(
                label="Load Example",
                choices=[""] + [ex["label"] for ex in EXAMPLES],
                value="",
                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")

            submit = gr.Button("Rerank", variant="primary")

        # ── 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()

    # ── Events ──
    add_btn.click(
        fn=add_doc,
        inputs=[doc_count],
        outputs=doc_textboxes + [doc_count],
    )
    remove_btn.click(
        fn=remove_doc,
        inputs=[doc_count],
        outputs=doc_textboxes + [doc_count],
    )
    example_selector.change(
        fn=load_example,
        inputs=[example_selector],
        outputs=[query_input] + doc_textboxes + [doc_count],
    )
    submit.click(
        fn=rerank,
        inputs=[query_input, instruction_input, chunk_size_input] + doc_textboxes,
        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()