Create app.py

app.py

@@ -0,0 +1,161 @@
+"""
+KaLM-Reranker-V1-Nano on Hugging Face ZeroGPU Space.
+
+Key design choices for ZeroGPU:
+  1. `kalm_reranker.py` is downloaded from the model repo at startup (small file, no GPU needed).
+  2. The `KaLMReranker` instance itself is created lazily INSIDE the `@spaces.GPU`
+     function, because at module-load time ZeroGPU has not yet allocated a CUDA device.
+  3. `@spaces.GPU(duration=120)` asks ZeroGPU for up to 120 s of GPU time per call —
+     enough for one reranking batch on a 0.27B model.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from typing import List, Tuple
+
+import gradio as gr
+import spaces
+from huggingface_hub import hf_hub_download
+
+# --------------------------------------------------------------------------- #
+# 1. Bring in the `kalm_reranker.py` helper from the model repo.
+#    This file is ~14 KB; downloading it once at startup is much cleaner than
+#    copy-pasting the source into the Space repo.
+# --------------------------------------------------------------------------- #
+MODEL_ID = "KaLM-Embedding/KaLM-Reranker-V1-Nano"
+
+_kalm_helper_path = hf_hub_download(repo_id=MODEL_ID, filename="kalm_reranker.py")
+sys.path.insert(0, os.path.dirname(_kalm_helper_path))
+
+from kalm_reranker import KaLMReranker  # noqa: E402  (path set at runtime)
+
+# --------------------------------------------------------------------------- #
+# 2. Lazy-loaded global reranker. Instantiated on first call inside @spaces.GPU.
+# --------------------------------------------------------------------------- #
+_reranker: KaLMReranker | None = None
+
+
+def get_reranker() -> KaLMReranker:
+    global _reranker
+    if _reranker is None:
+        print(f"[KaLM] Loading model {MODEL_ID} on GPU ...", flush=True)
+        _reranker = KaLMReranker(
+            MODEL_ID,
+            device="cuda",          # ZeroGPU makes CUDA visible inside @spaces.GPU
+            dtype="bfloat16",       # BF16 matches the published checkpoint
+            batch_size=32,
+            query_max_length=512,
+            max_length=1024,        # encoder tokens for "<Document>: {passage}"
+            chunk_size=4,           # encoder chunk pooling (Matryoshka)
+        )
+        print("[KaLM] Model ready.", flush=True)
+    return _reranker
+
+
+# --------------------------------------------------------------------------- #
+# 3. Inference function — wrapped with @spaces.GPU so ZeroGPU allocates a GPU.
+# --------------------------------------------------------------------------- #
+DEFAULT_QUERY = "What is the capital of China?"
+DEFAULT_DOCS = (
+    "The capital of China is Beijing.\n"
+    "Gravity attracts bodies toward one another.\n"
+    "Beijing is located in northern China and has a population of over 20 million."
+)
+DEFAULT_INSTRUCTION = "Given a query, retrieve documents that answer the query."
+
+
+@spaces.GPU(duration=120)
+def rerank(
+    query: str,
+    documents_text: str,
+    instruction: str,
+) -> Tuple[str, list]:
+    """Rerank documents against a query.
+
+    Args:
+        query: User query string.
+        documents_text: Newline-separated candidate documents.
+        instruction: Task instruction (e.g. "Given a query, retrieve ...").
+
+    Returns:
+        (formatted_text, raw_rankings_json)
+    """
+    query = (query or "").strip()
+    if not query:
+        return "ERROR: query is empty.", []
+
+    docs: List[str] = [d.strip() for d in (documents_text or "").splitlines() if d.strip()]
+    if not docs:
+        return "ERROR: please provide at least one document (one per line).", []
+
+    instruction = (instruction or "").strip() or None  # fall back to model default
+
+    r = get_reranker()
+    rankings = r.rank(query, docs, instruction=instruction)
+
+    # Pretty-print
+    lines = [f"Top {len(rankings)} results for query: \"{query}\""]
+    lines.append("-" * 60)
+    for i, item in enumerate(rankings, start=1):
+        cid = item["corpus_id"]
+        score = item["score"]
+        snippet = docs[cid][:100].replace("\n", " ")
+        lines.append(f"#{i:>2}  score={score:.6f}  doc#{cid + 1}  {snippet}{'...' if len(docs[cid]) > 100 else ''}")
+    return "\n".join(lines), [
+        {"rank": i, "corpus_id": item["corpus_id"], "score": float(item["score"]), "document": docs[item["corpus_id"]]}
+        for i, item in enumerate(rankings, start=1)
+    ]
+
+
+# --------------------------------------------------------------------------- #
+# 4. Gradio UI
+# --------------------------------------------------------------------------- #
+with gr.Blocks(title="KaLM-Reranker-V1-Nano", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        "# 🚀 KaLM-Reranker-V1-Nano · ZeroGPU Demo\n"
+        "Encoder-decoder reranker (0.27B activated, BF16) served on Hugging Face ZeroGPU."
+    )
+
+    with gr.Row():
+        query_in = gr.Textbox(
+            label="Query",
+            value=DEFAULT_QUERY,
+            scale=2,
+            placeholder="The user query ...",
+        )
+        instr_in = gr.Textbox(
+            label="Instruction",
+            value=DEFAULT_INSTRUCTION,
+            scale=3,
+            placeholder="e.g. Given a query, retrieve documents that answer the query.",
+        )
+
+    docs_in = gr.Textbox(
+        label="Candidate documents (one per line)",
+        value=DEFAULT_DOCS,
+        lines=8,
+        placeholder="Paste one document per line ...",
+    )
+
+    btn = gr.Button("Rerank", variant="primary")
+    out_text = gr.Textbox(label="Ranking", lines=10, show_copy_button=True)
+    out_json = gr.JSON(label="Raw scores (rank / corpus_id / score / document)")
+
+    btn.click(
+        fn=rerank,
+        inputs=[query_in, docs_in, instr_in],
+        outputs=[out_text, out_json],
+    )
+
+    gr.Markdown(
+        "---\n"
+        "ℹ️ **First call after a cold start** takes ~30–60 s (the model is downloaded and moved "
+        "onto a freshly-allocated ZeroGPU). Subsequent calls in the same session are fast.\n\n"
+        "Model: [`KaLM-Embedding/KaLM-Reranker-V1-Nano`](https://huggingface.co/KaLM-Embedding/KaLM-Reranker-V1-Nano)"
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()