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


def parse_documents(text: str) -> List[str]:
    docs = [doc.strip() for doc in text.split("\n\n") if doc.strip()]
    docs = docs[:MAX_DOCS]
    docs = [doc[:MAX_DOC_CHARS] for doc in docs]
    return docs


@spaces.GPU
def rerank(query: str, documents: str, instruction: str):
    reranker = KaLMReranker(
        MODEL_ID,
        device=None,
        dtype=None,
        batch_size=4,
        query_max_length=512,
        max_length=1024,
        chunk_size=4,
    )

    query = query.strip()
    instruction = instruction.strip() or DEFAULT_INSTRUCTION
    docs = parse_documents(documents)

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

    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}`. Higher score means more relevant."
        )
        return table, summary
    except Exception as error:
        return [], f"Error during reranking: {repr(error)}"


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.
**Document format:** separate candidate documents with one blank line.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            query = gr.Textbox(
                label="Query",
                value="What is the capital of China?",
                lines=3,
            )
            instruction = gr.Textbox(
                label="Instruction",
                value=DEFAULT_INSTRUCTION,
                lines=2,
            )
            documents = gr.Textbox(
                label="Candidate Documents",
                value=(
                    "The capital of China is Beijing.\n\n"
                    "Gravity attracts bodies toward one another.\n\n"
                    "Shanghai is a major city in China.\n\n"
                    "Paris is the capital of France."
                ),
                lines=14,
            )
            submit = gr.Button("Rerank", variant="primary")

        with gr.Column(scale=1):
            output_table = gr.Dataframe(
                headers=["Rank", "Corpus ID", "Score", "Document"],
                label="Reranking Results",
                wrap=True,
            )
            output_summary = gr.Markdown()

    submit.click(
        fn=rerank,
        inputs=[query, documents, instruction],
        outputs=[output_table, output_summary],
    )

    gr.Examples(
        examples=[
            [
                "What is the capital of China?",
                (
                    "The capital of China is Beijing.\n\n"
                    "Gravity attracts bodies toward one another.\n\n"
                    "Paris is the capital of France."
                ),
                DEFAULT_INSTRUCTION,
            ],
            [
                "Which model is suitable for efficient reranking?",
                (
                    "KaLM-Reranker-V1-Nano is designed for efficient reranking.\n\n"
                    "Large language models are often expensive for reranking.\n\n"
                    "Image classifiers are used for visual recognition."
                ),
                DEFAULT_INSTRUCTION,
            ],
            [
                "What is KaLM-Reranker-V1 designed for?",
                (
                    "KaLM-Reranker-V1 is a reranker for compressed document reranking.\n\n"
                    "KaLM-Embedding is a general-purpose embedding model.\n\n"
                    "Weather forecasting predicts future weather conditions."
                ),
                DEFAULT_INSTRUCTION,
            ],
        ],
        inputs=[query, documents, instruction],
    )

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