Update app.py

app.py

@@ -1,30 +1,26 @@
 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 load_model():
-    """
-    Load KaLM-Reranker once when the Space starts.
- 
-    device=None and dtype=None allow the official wrapper to choose:
-    - BF16 on CUDA
-    - FP32 on CPU
-    """
+
+
+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,
@@ -34,78 +30,49 @@ def load_model():
         max_length=1024,
         chunk_size=4,
     )
-    return reranker
- 
- 
-reranker = load_model()
- 
- 
-def parse_documents(text: str) -> List[str]:
-    """
-    Split documents by blank lines.
-    Each candidate document should be separated by one empty line.
-    """
-    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
- 
- 
-def rerank(query: str, documents: str, instruction: str):
+
     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,
-                ]
-            )
- 
+            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(
@@ -113,13 +80,11 @@ Input a query and several candidate documents. The demo returns relevance scores
                 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=(
@@ -130,9 +95,8 @@ Input a query and several candidate documents. The demo returns relevance scores
                 ),
                 lines=14,
             )
- 
             submit = gr.Button("Rerank", variant="primary")
- 
+
         with gr.Column(scale=1):
             output_table = gr.Dataframe(
                 headers=["Rank", "Corpus ID", "Score", "Document"],
@@ -140,13 +104,13 @@ Input a query and several candidate documents. The demo returns relevance scores
                 wrap=True,
             )
             output_summary = gr.Markdown()
- 
+
     submit.click(
         fn=rerank,
         inputs=[query, documents, instruction],
         outputs=[output_table, output_summary],
     )
- 
+
     gr.Examples(
         examples=[
             [
@@ -179,7 +143,7 @@ Input a query and several candidate documents. The demo returns relevance scores
         ],
         inputs=[query, documents, instruction],
     )
- 
+
     gr.Markdown(
         """
 ## Citation