import asyncio import json import os from typing import List import evaluate import jieba from llama_index.core.embeddings import BaseEmbedding from llama_index.core.evaluation import SemanticSimilarityEvaluator from llama_index.core.schema import NodeWithScore from pydantic import BaseModel from metagpt.const import EXAMPLE_BENCHMARK_PATH from metagpt.logs import logger from metagpt.rag.factories import get_rag_embedding from metagpt.utils.common import read_json_file class DatasetInfo(BaseModel): name: str document_files: List[str] gt_info: List[dict] class DatasetConfig(BaseModel): datasets: List[DatasetInfo] class RAGBenchmark: def __init__( self, embed_model: BaseEmbedding = None, ): self.evaluator = SemanticSimilarityEvaluator( embed_model=embed_model or get_rag_embedding(), ) def _set_metrics( self, bleu_avg :float = 0.0, bleu_1 :float = 0.0, bleu_2 :float = 0.0, bleu_3 :float = 0.0, bleu_4 :float = 0.0, rouge_l :float = 0.0, semantic_similarity :float = 0.0, recall :float = 0.0, hit_rate :float = 0.0, mrr :float = 0.0, length :float = 0.0, generated_text :str = None, ground_truth_text: str = None, question: str = None ): metrics = { "bleu-avg": bleu_avg, "bleu-1": bleu_1, "bleu-2": bleu_2, "bleu-3": bleu_3, "bleu-4": bleu_4, "rouge-L": rouge_l, "semantic similarity": semantic_similarity, "recall": recall, "hit_rate": hit_rate, "mrr": mrr, "length": length, } log = { "generated_text": generated_text, "ground_truth_text": ground_truth_text, "question": question, } return {"metrics": metrics, "log": log} def bleu_score(self, response: str, reference: str, with_penalty=False) -> float: f = lambda text: list(jieba.cut(text)) bleu = evaluate.load(path="bleu") results = bleu.compute(predictions=[response], references=[[reference]], tokenizer=f) bleu_avg = results["bleu"] bleu1 = results["precisions"][0] bleu2 = results["precisions"][1] bleu3 = results["precisions"][2] bleu4 = results["precisions"][3] brevity_penalty = results["brevity_penalty"] if with_penalty: return bleu_avg, bleu1, bleu2, bleu3, bleu4 else: return 0.0 if brevity_penalty == 0 else bleu_avg / brevity_penalty, bleu1, bleu2, bleu3, bleu4 def rougel_score(self, response: str, reference: str) -> float: # pip install rouge_score f = lambda text: list(jieba.cut(text)) rouge = evaluate.load(path="rouge") results = rouge.compute(predictions=[response], references=[[reference]], tokenizer=f, rouge_types=["rougeL"]) score = results["rougeL"] return score def recall(self, nodes: list[NodeWithScore], reference_docs: list[str]) -> float: if nodes: total_recall = sum(any(node.text in doc for node in nodes) for doc in reference_docs) return total_recall / len(reference_docs) else: return 0.0 def HitRate(self, nodes: list[NodeWithScore], reference_docs: list[str]) -> float: if nodes: return 1.0 if any(node.text in doc for doc in reference_docs for node in nodes) else 0.0 else: return 0.0 def MRR(self, nodes: list[NodeWithScore], reference_docs: list[str]) -> float: mrr_sum = 0.0 for i, doc in enumerate(reference_docs, start=1): for node in nodes: if node.text in doc: mrr_sum += 1.0 / i break return mrr_sum / len(reference_docs) if reference_docs else 0.0 async def SemanticSimilarity(self, response: str, reference: str) -> float: result = await self.evaluator.aevaluate( response=response, reference=reference, ) return result.score async def compute_metric( self, response: str = None, reference: str = None, nodes: list[NodeWithScore] = None, reference_doc: list[str] = None, question: str = None, ): recall = self.recall(nodes, reference_doc) bleu_avg, bleu1, bleu2, bleu3, bleu4 = self.bleu_score(response, reference) rouge_l = self.rougel_score(response, reference) hit_rate = self.HitRate(nodes, reference_doc) mrr = self.MRR(nodes, reference_doc) similarity = await self.SemanticSimilarity(response, reference) result = self._set_metrics( bleu_avg, bleu1, bleu2, bleu3, bleu4, rouge_l, similarity, recall, hit_rate, mrr, len(response), response, reference, question ) return result @staticmethod def load_dataset(ds_names: list[str] = ["all"]): infos = read_json_file(os.path.join(EXAMPLE_BENCHMARK_PATH, "dataset_info.json")) dataset_config = DatasetConfig( datasets=[ DatasetInfo( name=name, document_files=[ os.path.join(EXAMPLE_BENCHMARK_PATH, name, file) for file in info["document_file"] ], gt_info=read_json_file(os.path.join(EXAMPLE_BENCHMARK_PATH, name, info["gt_file"])), ) for dataset_info in infos for name, info in dataset_info.items() if name in ds_names or "all" in ds_names ] ) return dataset_config if __name__ == "__main__": benchmark = RAGBenchmark() answer = "是的，根据提供的信息，2023年7月20日，应急管理部和财政部确实联合发布了《因灾倒塌、损坏住房恢复重建救助工作规范》的通知。这份《规范》旨在进一步规范因灾倒塌、损坏住房的恢复重建救助相关工作。它明确了地方各级政府负责实施救助工作，应急管理部和财政部则负责统筹指导。地方财政应安排足够的资金，中央财政也会提供适当的补助。救助资金将通过专账管理，并采取特定的管理方式。救助对象是那些因自然灾害导致住房倒塌或损坏，并向政府提出申请且符合条件的受灾家庭。相关部门将组织调查统计救助对象信息，并建立档案。此外，《规范》还强调了资金发放的具体方式和公开透明的要求。" ground_truth = "“启明行动”是为了防控儿童青少年的近视问题，并发布了《防控儿童青少年近视核心知识十条》。" bleu_avg, bleu1, bleu2, bleu3, bleu4 = benchmark.bleu_score(answer, ground_truth) logger.info(f"bleu_avg = {bleu_avg}") logger.info(f"bleu1 = {bleu1}") logger.info(f"bleu2 = {bleu2}") logger.info(f"bleu3 = {bleu3}") logger.info(f"bleu4 = {bleu4}") rougeL_score = benchmark.rougel_score(answer, ground_truth) logger.info(f"rougeL_score = {rougeL_score}") similarity = asyncio.run(benchmark.SemanticSimilarity(answer, ground_truth)) logger.info(f"similarity = {similarity}")