run.py 1.2 KB

123456789101112131415161718192021222324252627282930313233343536
  1. from law_rag.retrieval import high_level_retrieval
  2. def law_rag_run(questions, with_score: bool = False):
  3. retriever = high_level_retrieval.HighLevelRetriever()
  4. results = retriever.search(questions)
  5. print(f"\n总共找到 {results['total_results']} 条相关法律条文")
  6. print(f"耗时: {results['time_taken']:.2f}秒")
  7. law_retrieval_result_list = []
  8. for i, item in enumerate(results["results"][:10]):
  9. print(f"\n[{i + 1}] 相似度: {item.get('original_score', 0):.4f}")
  10. print(f"来源: {item.get('source', '未知')}")
  11. if "purpose" in item:
  12. print(f"目的: {item['purpose']}")
  13. law_id = item['metadata'].get('law_id', '未知')
  14. content = item['text']
  15. print(f"法律ID: {law_id}")
  16. print(f"内容: {content}")
  17. row = {"law_id": law_id, "content": content}
  18. if with_score:
  19. row["score"] = item.get("original_score", 0)
  20. row["source"] = item.get("source", "未知")
  21. if "purpose" in item:
  22. row["purpose"] = item["purpose"]
  23. law_retrieval_result_list.append(row)
  24. return law_retrieval_result_list
  25. if __name__ == '__main__':
  26. result = law_rag_run("欠薪三月未还")
  27. print(str(result))