使用 ZeroEntropy Zerank-2 重新排序器設計高精度檢索與重新排序管道

In this tutorial, we use zeroentropy/zerank-2-reranker, a 4B Qwen3-based cross-encoder reranker, to improve retrieval quality. We start by setting up the runtime, loading the reranker, and understanding how it scores query-document pairs. Then, we move from simple pairwise scoring to a practical two-stage retrieve-and-rerank pipeline, where a fast bi-encoder first retrieves candidates and zerank-2 reranks them for better precision. We also evaluate the impact using NDCG@10 and test the reranker across finance, legal, and code examples to assess its performance in real-world search and ranking tasks. Copy CodeCopiedUse a different Browser!pip -q install -U "sentence-transformers>=3.0" "transformers>=4.51.0" accelerate import os, time, numpy as np, torch from sentence_transformers import CrossEncoder, SentenceTransformer, util os.environ["TOKENIZERS_PARALLELISM"] = "false" if torch.cuda.is_available(): device = "cuda" dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 print(f"GPU: {torch.cuda.get_device_name(0)} | dtype: {dtype}") else: device, dtype = "cpu", torch.float32 print("WARNING: no GPU detected. This 4B model will be very slow on CPU.") RERANKER_ID = "zeroentropy/zerank-2-reranker" print(f"\nLoading {RERANKER_ID} (~8GB on first run)...") reranker = CrossEncoder( RERANKER_ID, model_kwargs={"torch_dtype": dtype}, device=device, ) print("Reranker loaded.") def to_prob(logits): return (torch.as_tensor(logits, dtype=torch.float32) / 5).sigmoid() We begin by installing the required libraries and importing the main tools needed for reranking and retrieval. We check whether a GPU is available and select the appropriate device and tensor precision for efficient model execution. We then load the zeroentropy/zerank-2-reranker model and define a helper function to convert raw logits into probability-style scores. Copy CodeCopiedUse a different Browserprint("\n" + "="*70 + "\nPART 1: Pairwise scoring\n" + "="*70) pairs = [ ("What is 2+2?", "4"), ("What is 2+2?", "The answer is definitely 1 million"), ("Which planet is the Red Planet?", "Mars, known for its reddish appearance, is the Red Planet."), ("Which planet is the Red Planet?", "Venus is Earth's twin because of its similar size."), ] logits = reranker.predict(pairs, convert_to_tensor=True) probs = to_prob(logits) for (q, d), lg, p in zip(pairs, logits.tolist(), probs.tolist()): print(f"logit={lg:+6.2f} prob={p:5.3f} | {q[:30]:30s} -> {d[:45]}") ̧We test the reranker on simple query-document pairs to understand how it scores relevant and irrelevant answers. We pass each pair through reranker.predict() and receive raw logits from the model. We convert those logits into probabilities and print the results so we can compare how strongly the model prefers correct responses. Copy CodeCopiedUse a different Browserprint("\n" + "="*70 + "\nPART 2: model.rank for a single query\n" + "="*70) query = "How do I fix a Python list index out of range error?" candidates = [ "IndexError happens when you access an index beyond the list length; check len() and loop bounds.", "Use a try/except IndexError block, or validate the index with `if i < len(lst)` before access.", "To install Python packages, run `pip install <package>` in your terminal.", "List comprehensions create new lists: `[x*2 for x in nums]`.", "Off-by-one errors in range(len(lst)+1) are a common cause of index out of range.", ] ranking = reranker.rank(query, candidates, convert_to_tensor=True) for rank, r in enumerate(ranking, 1): cid = r["corpus_id"] print(f"#{rank} score={float(r['score']):+6.2f} prob={to_prob(r['score']):.3f} " f"| {candidates[cid][:60]}") We use model.rank() to rank multiple candidate answers for a single query. We provide several possible explanations for a Python list index error and let the reranker order them by relevance. We then print each ranked result with its raw score and probability-style score to see which answer the model considers most useful. Copy CodeCopiedUse a different Browserprint("\n" + "="*70 + "\nPART 3: Two-stage retrieve -> rerank pipeline\n" + "="*70) corpus = [ "The mitochondria is the powerhouse of the cell, producing ATP via respiration.", "Photosynthesis converts light energy into chemical energy in chloroplasts.", "ATP synthase uses a proton gradient across the inner mitochondrial membrane to make ATP.", "DNA replication is semi-conservative and occurs during the S phase of the cell cycle.", "The Krebs cycle (citric acid cycle) takes place in the mitochondrial matrix.", "Ribosomes translate mRNA into proteins in the cytoplasm.", "Glycolysis breaks glucose into pyruvate in the cytosol, yielding a net 2 ATP.", "The Golgi apparatus modifies, sorts, and packages proteins for secretion.", "Cellular respiration in mitochondria yields far more ATP than glycolysis alone.", "Plant cell walls are made primarily of cellulose for structural support.", ] bi = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device=device) corpus_emb = bi.encode(corpus, convert_to_tensor=True, normalize_embeddings=True) def two_stage_search(q, top_k_retrieve=6, top_n_final=3): q_emb = bi.encode(q, convert_to_tensor=True, normalize_embeddings=True) hits = util.semantic_search(q_emb, corpus_emb, top_k=top_k_retrieve)[0] cand_ids = [h["corpus_id"] for h in hits] cand_docs = [corpus[i] for i in cand_ids] rr = reranker.rank(q, cand_docs, convert_to_tensor=True) out = [] for r in rr[:top_n_final]: global_id = cand_ids[r["corpus_id"]] out.append((global_id, corpus[global_id], float(to_prob(r["score"])))) return cand_ids, out q = "Where in the cell is most ATP actually produced?" retrieved, final = two_stage_search(q) print(f"Query: {q}\n") print("Stage 1 (bi-encoder) top order:", retrieved) print("\nStage 2 (zerank-2 reranked) top results:") for gid, doc, p in final: print(f" [doc {gid}] prob={p:.3f} | {doc}") We build a two-stage retrieval pipeline that first uses a fast bi-encoder to retrieve candidate documents from a small corpus. We then pass those retrieved candidates to zerank-2 so it can rerank them with deeper query-document understanding. We finally compare the initially retrieved order with the reranked top results to see how reranking improves precision. Copy CodeCopiedUse a different Browserprint("\n" + "="*70 + "\nPART 4: NDCG@10 evaluation\n" + "="*70) eval_set = [ {"query": "Where is most ATP produced in the cell?", "rels": {0: 2, 2: 3, 4: 2, 6: 1, 8: 3}}, {"query": "How do plants capture light energy?", "rels": {1: 3, 9: 1}}, {"query": "How are proteins made and packaged in a cell?", "rels": {5: 3, 7: 2}}, ] def dcg(rels): rels = np.asarray(rels, dtype=float) return np.sum((2**rels - 1) / np.log2(np.arange(2, rels.size + 2))) def ndcg_at_k(ranked_doc_ids, rel_map, k=10): gains = [rel_map.get(d, 0) for d in ranked_doc_ids[:k]] ideal = sorted(rel_map.values(), reverse=True)[:k] idcg = dcg(ideal) return dcg(gains) / idcg if idcg > 0 else 0.0 base_scores, rr_scores = [], [] for ex in eval_set: q, rel_map = ex["query"], ex["rels"] q_emb = bi.encode(q, convert_to_tensor=True, normalize_embeddings=True) hits = util.semantic_search(q_emb, corpus_emb, top_k=len(corpus))[0] base_order = [h["corpus_id"] for h in hits] base_scores.append(ndcg_at_k(base_order, rel_map)) rr = reranker.rank(q, [corpus[i] for i in base_order], convert_to_tensor=True) rr_order = [base_order[r["corpus_id"]] for r in rr] rr_scores.append(ndcg_at_k(rr_order, rel_map)) print(f"{'Query':45s} {'bi-encoder':>12s} {'+ zerank-2':>12s}") for ex, b, r in zip(eval_set, base_scores, rr_scores): print(f"{ex['query'][:43]:45s} {b:12.4f} {r:12.4f}") print("-"*72) print(f"{'AVERAGE NDCG@10':45s} {np.mean(base_scores):12.4f} {np.mean(rr_scores):12.4f}") print(f"\nReranking lift: {np.mean(rr_scores)-np.mean(base_scores):+.4f} NDCG@10") We evaluate the retrieval pipeline using a small labeled benchmark and the NDCG@10 metric. We first measure the ranking quality of the bi-encoder alone and then measure the quality after applying zerank-2