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