diff --git a/rag2_0/tool/ModelTool.py b/rag2_0/tool/ModelTool.py index c3dfe78..f5f4bc2 100644 --- a/rag2_0/tool/ModelTool.py +++ b/rag2_0/tool/ModelTool.py @@ -176,8 +176,57 @@ class OpenAiLLM: return completion.choices[0].message if __name__ == "__main__": + # 测试重排模型 reranker = SiliconFlowReRankerModel() - query = "什么是AI" - documents = ["AI是人工智能", "AI是机器学习", "AI是深度学习"] + + # 测试用例1:简单问题 + query = "他想做什么" + documents = ["她想去公园跑步", "她想换一个新手机", "明天她想出去旅游"] results = reranker.rerank(query, documents) - print(results) + print(f"测试用例1 - 查询:{query}") + for idx, item in enumerate(results): + print(f"{idx+1}. 文档: {item['document']}, 分数: {item['score']}") + print("-" * 50) + + # 测试用例2:技术问题 + query = "Python如何处理JSON数据" + documents = [ + "Python中可以使用json模块来处理JSON数据,例如json.loads()将JSON字符串转换为字典", + "Java提供了多种处理JSON的库,比如Jackson和Gson", + "在Python中,可以使用pandas库来分析CSV数据", + "JavaScript可以使用JSON.parse()方法解析JSON字符串" + ] + results = reranker.rerank(query, documents) + print(f"测试用例2 - 查询:{query}") + for idx, item in enumerate(results): + print(f"{idx+1}. 文档: {item['document']}, 分数: {item['score']}") + print("-" * 50) + + # 测试用例3:医疗问题 + query = "高血压的症状有哪些" + documents = [ + "高血压的常见症状包括头痛、头晕、耳鸣和视力模糊", + "糖尿病的症状包括多饮、多尿和体重减轻", + "心脏病的症状通常包括胸痛、呼吸急促和疲劳", + "高血压患者应该定期监测血压,保持健康的生活方式" + ] + results = reranker.rerank(query, documents) + print(f"测试用例3 - 查询:{query}") + for idx, item in enumerate(results): + print(f"{idx+1}. 文档: {item['document']}, 分数: {item['score']}") + print("-" * 50) + + # 测试用例4:长文本查询和文档 + query = "人工智能在医疗领域的应用及其伦理问题" + documents = [ + "人工智能在医疗诊断中的应用已经显示出良好的效果,例如通过分析医学影像来检测疾病。然而,这也引发了关于医生角色和责任的伦理问题。", + "在教育领域,人工智能可以提供个性化学习体验,适应不同学生的学习进度和风格。", + "医疗伦理问题主要包括患者隐私保护、知情同意和医疗资源分配等方面。", + "人工智能技术在金融领域的应用主要集中在风险评估、欺诈检测和算法交易等方面。" + ] + results = reranker.rerank(query, documents) + print(f"测试用例4 - 查询:{query}") + for idx, item in enumerate(results): + print(f"{idx+1}. 文档: {item['document']}, 分数: {item['score']}") + print("-" * 50) +