diff --git a/docs/guides/Qwen3_Reranker.md b/docs/guides/Qwen3_Reranker.md index e082934d..76e03276 100644 --- a/docs/guides/Qwen3_Reranker.md +++ b/docs/guides/Qwen3_Reranker.md @@ -39,6 +39,16 @@ pip install onnxruntime # For ONNX inference ```python from embed_anything import Reranker, Dtype +# Qwen3 Reranker requires additional formatting +def format_query(query: str, instruction=None): + """You may add instruction to get better results in specific fields.""" + if instruction is None: + instruction = "Given a web search query, retrieve relevant passages that answer the query" + return f": {instruction}\n: {query}\n" + +def format_document(doc: str): + return f": {doc}" + # Initialize the Qwen3 reranker reranker = Reranker.from_pretrained( "zhiqing/Qwen3-Reranker-0.6B-ONNX", @@ -54,6 +64,10 @@ documents = [ "Pizza is a popular food." ] +# Format query and documents +query = [format_query(x) for x in query] +documents = [format_document(x) for x in documents] + results = reranker.rerank(query, documents, top_k=2) # Display results @@ -81,6 +95,10 @@ documents = [ "Python is great for beginners." ] +# Format queries and documents +queries = [format_query(x) for x in queries] +documents = [format_document(x) for x in documents] + results = reranker.rerank(queries, documents, top_k=3) for result in results: diff --git a/examples/qwen3_reranker.py b/examples/qwen3_reranker.py index b5b9975a..093fe2d8 100644 --- a/examples/qwen3_reranker.py +++ b/examples/qwen3_reranker.py @@ -14,17 +14,26 @@ from embed_anything import Reranker, Dtype, RerankerResult, DocumentRank import time +def format_query(query: str, instruction=None): + """You may add instruction to get better results in specific fields.""" + if instruction is None: + instruction = "Given a web search query, retrieve relevant passages that answer the query" + return f": {instruction}\n: {query}\n" + +def format_document(doc: str): + return f": {doc}" + def basic_qwen3_reranking(): """Basic example of using Qwen3 reranker for simple document ranking.""" print("=== Basic Qwen3 Reranking ===") - + # Initialize the Qwen3 reranker # Using the ONNX-optimized version for better performance reranker = Reranker.from_pretrained( "zhiqing/Qwen3-Reranker-0.6B-ONNX", dtype=Dtype.F32 ) - + # Define a query and candidate documents query = ["What is artificial intelligence?"] documents = [ @@ -35,10 +44,14 @@ def basic_qwen3_reranking(): "Deep learning uses neural networks.", "Pizza is a popular Italian food." ] - + + # Format query and documents + query = [format_query(x) for x in query] + documents = [format_document(x) for x in documents] + # Rerank documents and get top 3 results results = reranker.rerank(query, documents, 2) - + # Display results for result in results: print(f"Query: {result.query}") @@ -51,19 +64,19 @@ def basic_qwen3_reranking(): def multi_query_reranking(): """Example of reranking documents for multiple queries simultaneously.""" print("=== Multi-Query Reranking ===") - + reranker = Reranker.from_pretrained( "zhiqing/Qwen3-Reranker-0.6B-ONNX", dtype=Dtype.F32 ) - + # Multiple queries queries = [ "How to make coffee?", "What is machine learning?", "Tell me about cats" ] - + # Shared document collection documents = [ "Coffee is made by brewing ground coffee beans with hot water.", @@ -77,10 +90,14 @@ def multi_query_reranking(): "Coffee beans come from the Coffea plant.", "Cats are known for their independent nature." ] - + + # Format queries and documents + queries = [format_query(x) for x in queries] + documents = [format_document(x) for x in documents] + # Rerank for all queries at once results = reranker.rerank(queries, documents, top_k=3) - + # Display results for each query for result in results: print(f"Query: {result.query}") @@ -92,12 +109,12 @@ def multi_query_reranking(): def custom_scoring_example(): """Example of using compute_scores for custom ranking logic.""" print("=== Custom Scoring with compute_scores ===") - + reranker = Reranker.from_pretrained( "zhiqing/Qwen3-Reranker-0.6B-ONNX", dtype=Dtype.F32 ) - + query = ["What are the benefits of exercise?"] documents = [ "Exercise improves cardiovascular health.", @@ -108,31 +125,35 @@ def custom_scoring_example(): "Physical activity increases energy levels.", "Exercise promotes better sleep quality." ] - + + # Format query and documents + query = [format_query(x) for x in query] + documents = [format_document(x) for x in documents] + # Get raw scores for custom processing scores = reranker.compute_scores(query, documents, batch_size=4) - + print(f"Raw relevance scores for query: '{query[0]}'") print("-" * 50) - + # Create custom ranking based on scores doc_scores = list(zip(documents, scores[0])) doc_scores.sort(key=lambda x: x[1], reverse=True) - + for i, (doc, score) in enumerate(doc_scores): print(f"{i+1:2d}. Score: {score:6.4f} | {doc}") - + print() def performance_benchmark(): """Benchmark the performance of Qwen3 reranking.""" print("=== Performance Benchmark ===") - + reranker = Reranker.from_pretrained( "zhiqing/Qwen3-Reranker-0.6B-ONNX", dtype=Dtype.F32 ) - + # Create a larger dataset for benchmarking queries = ["What is technology?"] * 5 # 5 identical queries documents = [ @@ -147,20 +168,24 @@ def performance_benchmark(): "Educational technology enhances learning.", "Space technology enables exploration." ] * 10 # 100 total documents - + + # Format queries and documents + queries = [format_query(x) for x in queries] + documents = [format_document(x) for x in documents] + print(f"Benchmarking with {len(queries)} queries and {len(documents)} documents...") - + # Warm up _ = reranker.compute_scores(queries[:1], documents[:10], batch_size=4) - + # Benchmark start_time = time.time() results = reranker.rerank(queries, documents, top_k=5) end_time = time.time() - + processing_time = end_time - start_time docs_per_second = len(documents) / processing_time - + print(f"Processing time: {processing_time:.2f} seconds") print(f"Documents processed per second: {docs_per_second:.1f}") print(f"Total documents processed: {len(documents)}") @@ -169,15 +194,15 @@ def performance_benchmark(): def search_and_rerank_pipeline(): """Example of a complete search and rerank pipeline.""" print("=== Search and Rerank Pipeline ===") - + reranker = Reranker.from_pretrained( "zhiqing/Qwen3-Reranker-0.6B-ONNX", dtype=Dtype.F32 ) - + # Simulate a search query search_query = ["How to learn Python programming?"] - + # Simulate candidate documents from a vector search candidate_docs = [ "Python is a high-level programming language.", @@ -193,17 +218,21 @@ def search_and_rerank_pipeline(): "Start with simple projects when learning.", "Python has a large and active community." ] - + print(f"Search Query: {search_query[0]}") print(f"Found {len(candidate_docs)} candidate documents") print("\nReranking documents by relevance...") - + + # Format search_query and documents + search_query = [format_query(x) for x in search_query] + documents = [format_document(x) for x in documents] + # Rerank the candidates reranked_results = reranker.rerank(search_query, candidate_docs, top_k=5) - + print("\nTop 5 most relevant documents:") print("-" * 60) - + for result in reranked_results: for doc in result.documents: print(f"Rank {doc.rank:2d} (Score: {doc.relevance_score:.4f}):") @@ -214,7 +243,7 @@ def search_and_rerank_pipeline(): print("Qwen3 Reranker Examples") print("=" * 50) print() - + try: # Run all examples basic_qwen3_reranking() @@ -222,9 +251,9 @@ def search_and_rerank_pipeline(): custom_scoring_example() performance_benchmark() search_and_rerank_pipeline() - + print("All examples completed successfully!") - + except Exception as e: print(f"Error running examples: {e}") print("Make sure you have the required dependencies installed:") diff --git a/rust/src/reranker/model.rs b/rust/src/reranker/model.rs index 58ca6fa2..d6a4e731 100644 --- a/rust/src/reranker/model.rs +++ b/rust/src/reranker/model.rs @@ -155,7 +155,7 @@ impl Reranker { // Check model type once at the beginning let is_qwen3 = self.model_type.as_ref().is_some_and(|t| t == "qwen3"); - let prefix = "<|im_start|>system\nJudge 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"; + let prefix = "<|im_start|>system\nJudge 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"; let suffix = "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; let pairs = queries