Upload folder using huggingface_hub

README.md

@@ -1,28 +1,16 @@
 ---
-title: Text Reranker - Cross-Encoder Reranking
+title: Qwen3-Reranker-8B Text Reranker
 emoji: 🔎
 colorFrom: blue
 colorTo: indigo
 sdk: gradio
-sdk_version: 5.17.1
 app_file: app.py
 pinned: false
 license: mit
 ---
 
-# Text Reranker - Cross-Encoder Reranking
+# Qwen3-Reranker-8B Text Reranker
 
-A lightweight Hugging Face Space that reranks passages for a given query using:
+Fast text-only reranking Space powered by `Qwen/Qwen3-Reranker-8B`.
 
-- `cross-encoder/ms-marco-MiniLM-L-12-v2`
-- `sentence-transformers` `CrossEncoder`
-- Gradio UI with ZeroGPU support via `@spaces.GPU`
-
-## Usage
-
-1. Enter a query.
-2. Paste passages, one per line.
-3. Choose Top-K.
-4. Click **Rerank**.
-
-The app returns a markdown table sorted by relevance score and displays inference time.
+Enter a query and passages, one per line. The app returns a sorted relevance table and inference time.

app.py

@@ -3,10 +3,36 @@ from typing import List
 
 import gradio as gr
 import spaces
-from sentence_transformers import CrossEncoder
-
-MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-12-v2"
-model = CrossEncoder(MODEL_NAME)
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_NAME = "Qwen/Qwen3-Reranker-8B"
+INSTRUCTION = "Given a web search query, retrieve relevant passages that answer the query"
+
+# Load once at startup
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    trust_remote_code=True,
+)
+if torch.cuda.is_available():
+    model = model.cuda()
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="left", trust_remote_code=True)
+token_false_id = tokenizer.convert_tokens_to_ids("no")
+token_true_id = tokenizer.convert_tokens_to_ids("yes")
+
+max_length = 8192
+prefix = (
+    "<|im_start|>system\n"
+    "Judge whether the Document meets the requirements based on the Query and the Instruct provided. "
+    "Note that the answer can only be \"yes\" or \"no\"."
+    "<|im_end|>\n<|im_start|>user\n"
+)
+suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
+prefix_tokens = tokenizer.encode(prefix, add_special_tokens=False)
+suffix_tokens = tokenizer.encode(suffix, add_special_tokens=False)
 
 
 def _parse_passages(text: str) -> List[str]:
@@ -15,6 +41,36 @@ def _parse_passages(text: str) -> List[str]:
     return [line.strip() for line in text.splitlines() if line.strip()]
 
 
+def _format_pair(query: str, doc: str) -> str:
+    return f"<Instruct>: {INSTRUCTION}\n<Query>: {query}\n<Document>: {doc}"
+
+
+def _process_inputs(pairs: List[str]):
+    inputs = tokenizer(
+        pairs,
+        padding=False,
+        truncation="longest_first",
+        return_attention_mask=False,
+        max_length=max_length - len(prefix_tokens) - len(suffix_tokens),
+    )
+    for i, ids in enumerate(inputs["input_ids"]):
+        inputs["input_ids"][i] = prefix_tokens + ids + suffix_tokens
+    inputs = tokenizer.pad(inputs, padding=True, return_tensors="pt", max_length=max_length)
+    for key in inputs:
+        inputs[key] = inputs[key].to(model.device)
+    return inputs
+
+
+@torch.no_grad()
+def _compute_scores(inputs):
+    logits = model(**inputs).logits[:, -1, :]
+    true_vector = logits[:, token_true_id]
+    false_vector = logits[:, token_false_id]
+    score_2way = torch.stack([false_vector, true_vector], dim=1)
+    score_2way = torch.nn.functional.log_softmax(score_2way, dim=1)
+    return score_2way[:, 1].exp().tolist()
+
+
 @spaces.GPU
 def rerank(query: str, passages_text: str, top_k: int):
     start = time.perf_counter()
@@ -23,79 +79,44 @@ def rerank(query: str, passages_text: str, top_k: int):
     passages = _parse_passages(passages_text or "")
 
     if not query and not passages:
-        return (
-            "Please provide a query and at least one passage.",
-            "Inference time: 0.000s",
-        )
+        return "Please provide a query and at least one passage.", "Inference time: 0.000s"
     if not query:
         return "Please provide a query.", "Inference time: 0.000s"
     if not passages:
         return "Please provide at least one passage.", "Inference time: 0.000s"
 
-    top_k = max(1, min(int(top_k), 20, len(passages)))
+    top_k = max(1, min(int(top_k), 50, len(passages)))
+    pairs = [_format_pair(query, p) for p in passages]
+    inputs = _process_inputs(pairs)
+    scores = _compute_scores(inputs)
 
-    pairs = [[query, p] for p in passages]
-    scores = model.predict(pairs)
-
-    ranked = sorted(
-        zip(passages, scores),
-        key=lambda x: float(x[1]),
-        reverse=True,
-    )
+    ranked = sorted(zip(passages, scores), key=lambda x: float(x[1]), reverse=True)[:top_k]
 
-    ranked = ranked[:top_k]
-
-    lines = [
-        "| Rank | Score | Passage |",
-        "|---:|---:|---|",
-    ]
+    lines = ["| Rank | Score | Passage |", "|---:|---:|---|"]
     for i, (passage, score) in enumerate(ranked, start=1):
         safe_passage = passage.replace("|", "\\|").replace("\n", " ")
-        lines.append(f"| {i} | {float(score):.4f} | {safe_passage} |")
+        lines.append(f"| {i} | {float(score):.6f} | {safe_passage} |")
 
     elapsed = time.perf_counter() - start
     return "\n".join(lines), f"Inference time: {elapsed:.3f}s"
 
 
-with gr.Blocks(title="Text Reranker - Cross-Encoder Reranking") as demo:
-    gr.Markdown("# Text Reranker - Cross-Encoder Reranking")
-
-    query = gr.Textbox(
-        label="Query",
-        placeholder="Enter your search query...",
-        lines=1,
-    )
+with gr.Blocks(title="Qwen3-Reranker-8B Text Reranker") as demo:
+    gr.Markdown("# Qwen3-Reranker-8B Text Reranker")
 
+    query = gr.Textbox(label="Query", placeholder="Enter your search query...", lines=1)
     passages = gr.Textbox(
         label="Passages (one per line)",
         placeholder="Enter one passage per line...",
         lines=10,
     )
-
-    top_k = gr.Slider(
-        minimum=1,
-        maximum=20,
-        value=5,
-        step=1,
-        label="Top-K",
-    )
-
+    top_k = gr.Slider(minimum=1, maximum=50, value=5, step=1, label="Top-K")
     run_btn = gr.Button("Rerank")
 
     output_md = gr.Markdown(label="Ranked Results")
     inference_time = gr.Textbox(label="Inference Time", interactive=False)
 
-    run_btn.click(
-        fn=rerank,
-        inputs=[query, passages, top_k],
-        outputs=[output_md, inference_time],
-    )
-
-    gr.Markdown(
-        "Built by [Xavier Fuentes](https://huggingface.co/xavier-fuentes) @ "
-        "[AI Enablement Academy](https://enablement.academy) | "
-        "[Buy me a coffee ☕](https://ko-fi.com/xavierfuentes)"
-    )
+    run_btn.click(fn=rerank, inputs=[query, passages, top_k], outputs=[output_md, inference_time])
 
 
 if __name__ == "__main__":

requirements.txt

@@ -1,4 +1,5 @@
-gradio
-sentence-transformers
-torch
+transformers>=4.57.0
+torch>=2.0
+gradio>=4.0
+spaces
 accelerate