Lychee-Rerank代码实例：Python调用API实现批量文档重排序

Lychee-Rerank代码实例：Python调用API实现批量文档重排序

1. 项目简介

Lychee-Rerank是一个基于本地推理的文档相关性评分工具，它采用了Lychee官方推理逻辑并适配Qwen2.5-1.5B模型作为推理基座。这个工具专门为「查询-文档」匹配度打分场景设计，能够帮助开发者快速评估文档与查询语句的相关性。

在实际应用中，当你有一组候选文档和一个查询问题时，Lychee-Rerank可以自动为每个文档计算相关性分数，并按分数高低进行排序。整个过程完全在本地运行，不需要网络连接，既保护了数据隐私，又不受使用次数限制。

工具的核心特点是：

• 支持自定义评分指令和查询语句
• 批量处理多个候选文档
• 提供可视化进度条和颜色分级显示
• 纯本地推理，无需上传数据

2. 环境准备与安装

2.1 系统要求

在使用Lychee-Rerank之前，确保你的系统满足以下要求：

• Python 3.8或更高版本
• 至少8GB内存（处理大量文档时建议16GB以上）
• 支持CUDA的GPU（可选，但能显著加速推理过程）

2.2 安装依赖包

首先创建并激活Python虚拟环境：

python -m venv lychee-env source lychee-env/bin/activate # Linux/Mac # 或者 lychee-env\Scripts\activate # Windows

然后安装必要的依赖包：

pip install torch transformers streamlit tqdm

这些包的作用分别是：

• torch: 深度学习框架
• transformers: 预训练模型加载和推理
• streamlit: 构建可视化界面
• tqdm: 显示进度条

3. 核心代码实现

3.1 模型加载与初始化

下面是加载Qwen2.5-1.5B模型的核心代码：

import torch from transformers import AutoTokenizer, AutoModelForCausalLM import streamlit as st class LycheeReranker: def __init__(self, model_name="Qwen/Qwen2.5-1.5B"): self.device = "cuda" if torch.cuda.is_available() else "cpu" self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float16, device_map="auto" ) self.system_prompt = "判断以下文档是否与查询相关。只回答'是'或'否'。"

3.2 相关性评分计算

这是计算文档相关性分数的核心方法：

def calculate_relevance(self, instruction, query, document): # 构建完整的输入文本 full_text = f"<Instruction>{instruction}</Instruction>\n" full_text += f"<Query>{query}</Query>\n" full_text += f"<Document>{document}</Document>\n" full_text += "相关吗？回答是或否：" # 编码输入文本 inputs = self.tokenizer(full_text, return_tensors="pt").to(self.device) # 生成回答 with torch.no_grad(): outputs = self.model.generate( **inputs, max_new_tokens=10, return_dict_in_generate=True, output_scores=True ) # 计算"是"的概率作为相关性分数 yes_token_id = self.tokenizer.convert_tokens_to_ids("是") scores = torch.softmax(outputs.scores[0], dim=-1) relevance_score = scores[0, yes_token_id].item() return relevance_score

3.3 批量处理与排序

处理多个文档并排序的实现：

def batch_rerank(self, instruction, query, documents): results = [] # 为每个文档计算分数 for doc in documents: score = self.calculate_relevance(instruction, query, doc) results.append({ "document": doc, "score": score }) # 按分数降序排序 results.sort(key=lambda x: x["score"], reverse=True) return results

4. 可视化界面搭建

4.1 Streamlit界面配置

使用Streamlit构建用户友好的界面：

def main(): st.set_page_config(page_title="Lychee-Rerank 文档重排序", layout="wide") st.title("⚖ Lychee-Rerank 文档相关性评分工具") # 初始化模型 if "reranker" not in st.session_state: with st.spinner("正在加载模型，请稍候..."): st.session_state.reranker = LycheeReranker()

4.2 输入区域设计

创建用户输入区域：

# 创建两列布局 col1, col2 = st.columns([1, 2]) with col1: st.header("输入配置") # 指令输入 instruction = st.text_area( "评分指令", value="基于查询检索相关文档", height=100 ) # 查询输入 query = st.text_input("查询语句", "What is the capital of China?") # 文档输入 documents_text = st.text_area( "候选文档（每行一个）", height=300, value="北京是中国的首都。\n上海是中国最大的城市。\n巴黎是法国的首都。\n东京是日本的首都。\n伦敦是英国的首都。" ) documents = [doc.strip() for doc in documents_text.split("\n") if doc.strip()] # 计算按钮 if st.button(" 计算相关性分数", type="primary"): with st.spinner("计算中..."): results = st.session_state.reranker.batch_rerank(instruction, query, documents) st.session_state.results = results

4.3 结果展示优化

美化结果展示界面：

with col2: st.header("评分结果") if "results" in st.session_state: results = st.session_state.results for i, result in enumerate(results): score = result["score"] # 根据分数设置颜色 if score > 0.8: color = "green" relevance_level = "高相关性" elif score > 0.4: color = "orange" relevance_level = "中相关性" else: color = "red" relevance_level = "低相关性" # 显示每个文档的评分结果 st.markdown(f"### {i+1}. 得分: <span style='color:{color}'>{score:.6f}</span> ({relevance_level})", unsafe_allow_html=True) # 进度条可视化 st.progress(score) # 文档内容展示 st.code(result["document"], language="text")

5. 完整使用示例

5.1 基本使用流程

让我们通过一个完整例子来演示如何使用Lychee-Rerank：

1. 准备输入数据：

   • 指令：基于查询检索相关文档
   • 查询：人工智能的发展历史
   • 候选文档：5篇关于不同技术领域的文章

2. 运行评分：

   • 点击"计算相关性分数"按钮
   • 等待模型处理所有文档

3. 分析结果：

   • 查看按相关性排序的文档列表
   • 通过颜色标识快速识别高相关文档
   • 使用进度条直观比较相关性强度

5.2 实际应用场景

Lychee-Rerank在多个场景中都非常有用：

搜索引擎优化：对搜索结果进行重排序，提升用户体验

# 搜索引擎结果重排序示例 search_results = get_search_results("Python编程教程") reranked_results = reranker.batch_rerank( "找出最适合初学者的教程", "Python编程教程", search_results )

文档检索系统：从大量文档中快速找到最相关的内容

# 文档检索示例 all_documents = load_all_documents() relevant_docs = reranker.batch_rerank( "找出与查询最相关的技术文档", "如何安装PyTorch", all_documents )

内容推荐：根据用户查询推荐相关内容

# 内容推荐示例 user_query = "机器学习基础知识" content_library = get_all_articles() recommendations = reranker.batch_rerank( "推荐适合初学者的内容", user_query, content_library )

6. 高级功能与定制

6.1 自定义评分指令

你可以根据需要修改评分指令来适应不同的场景：

# 学术论文相关性评估 academic_instruction = "判断这篇论文是否与我的研究课题相关。考虑方法、数据集和结论的匹配度。" # 产品评论情感分析 sentiment_instruction = "判断这条评论是否表达了对产品功能的正面评价。" # 技术支持问答匹配 support_instruction = "判断这个解答是否准确回答了用户的技术问题。"

6.2 性能优化技巧

处理大量文档时，可以考虑以下优化措施：

批量处理优化：

def optimized_batch_rerank(self, instruction, query, documents, batch_size=4): results = [] # 分批处理避免内存溢出 for i in range(0, len(documents), batch_size): batch_docs = documents[i:i+batch_size] # 批量编码 batch_inputs = self.prepare_batch_inputs(instruction, query, batch_docs) # 批量推理 batch_scores = self.batch_calculate_relevance(batch_inputs) for doc, score in zip(batch_docs, batch_scores): results.append({"document": doc, "score": score}) results.sort(key=lambda x: x["score"], reverse=True) return results

模型量化加速：

# 使用8位量化减少内存使用 model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float16, device_map="auto", load_in_8bit=True # 8位量化 )

7. 常见问题与解决方案

7.1 内存不足问题

问题：处理大量文档时出现内存不足错误

解决方案：

# 减少批量大小 results = reranker.batch_rerank(instruction, query, documents, batch_size=2) # 使用内存映射 model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True # 减少CPU内存使用 )

7.2 评分一致性保障

问题：相同文档每次评分结果略有差异

解决方案：

# 设置随机种子确保可重复性 import random import numpy as np def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) set_seed(42) # 在推理前设置随机种子

7.3 处理长文档策略

问题：文档长度超过模型限制

解决方案：

def process_long_document(document, max_length=1000): # 简单截断策略 if len(document) > max_length: return document[:max_length] + "..." return document # 或者使用更智能的摘要策略 def summarize_long_document(document, max_length=1000): if len(document) <= max_length: return document # 这里可以添加更复杂的摘要逻辑 # 例如提取首段、尾段和关键句子 sentences = document.split('.') important_sentences = sentences[0:1] + sentences[-2:-1] return '.'.join(important_sentences) + '.'

8. 总结

Lychee-Rerank是一个强大且易用的文档相关性评分工具，它让复杂的重排序任务变得简单直观。通过本文的代码实例，你已经学会了如何：

1. 搭建基础环境：安装必要的依赖包和配置运行环境
2. 实现核心功能：加载模型、计算相关性分数、批量处理文档
3. 构建用户界面：使用Streamlit创建直观的可视化界面
4. 处理实际场景：应用在各种文档检索和排序任务中
5. 优化性能：通过批量处理和模型量化提升效率

这个工具的优势在于完全本地运行，保护数据隐私的同时提供专业级的文档重排序能力。无论是构建搜索引擎、文档检索系统，还是内容推荐引擎，Lychee-Rerank都能提供有力的支持。

现在你可以根据自己的需求，调整评分指令、优化处理流程，或者集成到更大的应用系统中，享受本地化文档重排序带来的便利和安全性。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。