first commit
This commit is contained in:
闫旭隆
1
AIEC-RAG/atlas_rag/evaluation/__init__.py
Normal file
1
AIEC-RAG/atlas_rag/evaluation/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
from .benchmark import BenchMarkConfig, RAGBenchmark
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
236
AIEC-RAG/atlas_rag/evaluation/benchmark.py
Normal file
236
AIEC-RAG/atlas_rag/evaluation/benchmark.py
Normal file
@ -0,0 +1,236 @@
|
||||
import os
|
||||
import json
|
||||
import numpy as np
|
||||
from logging import Logger
|
||||
from atlas_rag.retriever.base import BaseRetriever, BaseEdgeRetriever, BasePassageRetriever
|
||||
from typing import List
|
||||
from datetime import datetime
|
||||
from transformers import AutoModel
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from atlas_rag.vectorstore.embedding_model import NvEmbed, SentenceEmbedding
|
||||
from atlas_rag.llm_generator.llm_generator import LLMGenerator
|
||||
from atlas_rag.evaluation.evaluation import QAJudger
|
||||
from dataclasses import dataclass
|
||||
from atlas_rag.llm_generator.prompt.react import ReAct
|
||||
|
||||
|
||||
def normalize_embeddings(embeddings):
|
||||
"""Normalize the embeddings to unit length (L2 norm)."""
|
||||
if isinstance(embeddings, torch.Tensor):
|
||||
# Handle PyTorch tensors
|
||||
norm_emb = F.normalize(embeddings, p=2, dim=1).detach().cpu().numpy()
|
||||
elif isinstance(embeddings, np.ndarray):
|
||||
# Handle numpy arrays
|
||||
norm_emb = F.normalize(torch.tensor(embeddings), p=2, dim=1).detach().cpu().numpy()
|
||||
else:
|
||||
raise TypeError(f"Unsupported input type: {type(embeddings)}. Must be torch.Tensor or np.ndarray")
|
||||
|
||||
return norm_emb
|
||||
|
||||
@dataclass
|
||||
class BenchMarkConfig:
|
||||
"""
|
||||
Configuration class for benchmarking.
|
||||
|
||||
Attributes:
|
||||
dataset_name (str): Name of the dataset. Default is "hotpotqa".
|
||||
question_file (str): Path to the question file. Default is "hotpotqa".
|
||||
graph_file (str): Path to the graph file. Default is "hotpotqa_concept.graphml".
|
||||
include_events (bool): Whether to include events. Default is False.
|
||||
include_concept (bool): Whether to include concepts. Default is False.
|
||||
reader_model_name (str): Name of the reader model. Default is "meta-llama/Llama-2-7b-chat-hf".
|
||||
encoder_model_name (str): Name of the encoder model. Default is "nvidia/NV-Embed-v2".
|
||||
number_of_samples (int): Number of samples to use from the dataset. Default is -1 (use all samples).
|
||||
"""
|
||||
dataset_name: str = "hotpotqa"
|
||||
question_file: str = "hotpotqa"
|
||||
include_events: bool = False
|
||||
include_concept: bool = False
|
||||
reader_model_name: str = "meta-llama/Llama-2-7b-chat-hf"
|
||||
encoder_model_name: str = "nvidia/NV-Embed-v2"
|
||||
number_of_samples: int = -1 # Default to -1 to use all samples
|
||||
react_max_iterations: int = 5
|
||||
|
||||
|
||||
class RAGBenchmark:
|
||||
def __init__(self, config:BenchMarkConfig, logger:Logger = None):
|
||||
self.config = config
|
||||
self.logger = logger
|
||||
self.logging : bool = self.logger is not None
|
||||
|
||||
def load_encoder_model(self, encoder_model_name, **kwargs):
|
||||
if encoder_model_name == "nvidia/NV-Embed-v2":
|
||||
sentence_encoder = AutoModel.from_pretrained("nvidia/NV-Embed-v2", **kwargs)
|
||||
return NvEmbed(sentence_encoder)
|
||||
else:
|
||||
sentence_encoder = SentenceTransformer(encoder_model_name, **kwargs)
|
||||
return SentenceEmbedding(sentence_encoder)
|
||||
|
||||
def run(self, retrievers:List[BaseRetriever],
|
||||
llm_generator:LLMGenerator,
|
||||
use_react: bool = False):
|
||||
qa_judge = QAJudger()
|
||||
if use_react:
|
||||
react_agent = ReAct(llm_generator=llm_generator)
|
||||
result_list = []
|
||||
with open(self.config.question_file, "r") as f:
|
||||
data = json.load(f)
|
||||
print(f"Data loaded from {self.config.question_file}")
|
||||
if self.config.number_of_samples > 0:
|
||||
data = data[:self.config.number_of_samples]
|
||||
print(f"Using only the first {self.config.number_of_samples} samples from the dataset")
|
||||
for sample in tqdm(data):
|
||||
question = sample["question"]
|
||||
answer = sample["answer"]
|
||||
|
||||
gold_file_ids = []
|
||||
if self.config.dataset_name in ("hotpotqa", "2wikimultihopqa"):
|
||||
for fact in sample["supporting_facts"]:
|
||||
gold_file_ids.append(fact[0])
|
||||
elif self.config.dataset_name == "musique":
|
||||
for paragraph in sample["paragraphs"]:
|
||||
if paragraph["is_supporting"]:
|
||||
gold_file_ids.append(paragraph["paragraph_text"])
|
||||
else:
|
||||
print("Dataset not supported")
|
||||
continue
|
||||
|
||||
result = {
|
||||
"question": question,
|
||||
"answer": answer,
|
||||
"gold_file_ids": gold_file_ids,
|
||||
}
|
||||
|
||||
if self.logging:
|
||||
self.logger.info(f"Question: {question}")
|
||||
for retriever in retrievers:
|
||||
if use_react:
|
||||
# Use RAG with ReAct
|
||||
llm_generated_answer, search_history = react_agent.generate_with_rag_react(
|
||||
question=question,
|
||||
retriever=retriever,
|
||||
max_iterations=self.config.react_max_iterations,
|
||||
max_new_tokens=2048,
|
||||
logger=self.logger
|
||||
)
|
||||
self.logger.info(f"Search history: {search_history}")
|
||||
self.logger.info(f"Final answer: {llm_generated_answer}")
|
||||
# Store search history in results
|
||||
result[f"{retriever.__class__.__name__}_search_history"] = search_history
|
||||
|
||||
# Extract all retrieved contexts from search history
|
||||
all_contexts = []
|
||||
for _, action, observation in search_history:
|
||||
if "search" in action.lower() or "look up" in action.lower():
|
||||
all_contexts.append(observation)
|
||||
|
||||
sorted_context = "\n".join(all_contexts)
|
||||
sorted_context_ids = [] # We don't track IDs in ReAct mode
|
||||
else:
|
||||
# Original RAG implementation
|
||||
sorted_context, sorted_context_ids = retriever.retrieve(question, topN=5)
|
||||
|
||||
if isinstance(retriever, BaseEdgeRetriever):
|
||||
retrieved_context = ". ".join(sorted_context)
|
||||
llm_generated_answer = llm_generator.generate_with_context_kg(question, retrieved_context, max_new_tokens=2048, temperature=0.5)
|
||||
elif isinstance(retriever, BasePassageRetriever):
|
||||
retrieved_context = "\n".join(sorted_context)
|
||||
llm_generated_answer = llm_generator.generate_with_context(question, retrieved_context, max_new_tokens=2048, temperature=0.5)
|
||||
|
||||
if self.logging:
|
||||
self.logger.info(f"{retriever.__class__.__name__} retrieved passages: {sorted_context}")
|
||||
self.logger.info(f"{retriever.__class__.__name__} generated answer: {llm_generated_answer}")
|
||||
|
||||
short_answer = qa_judge.split_answer(llm_generated_answer)
|
||||
em, f1 = qa_judge.judge(short_answer, answer)
|
||||
|
||||
result[f"{retriever.__class__.__name__ }_em"] = em
|
||||
result[f"{retriever.__class__.__name__ }_f1"] = f1
|
||||
result[f"{retriever.__class__.__name__ }_passages"] = sorted_context
|
||||
if not use_react:
|
||||
result[f"{retriever.__class__.__name__ }_id"] = sorted_context_ids
|
||||
result[f"{retriever.__class__.__name__ }_generated_answer"] = llm_generated_answer
|
||||
result[f"{retriever.__class__.__name__ }short_answer"] = short_answer
|
||||
|
||||
# Calculate recall
|
||||
if not use_react: # Only calculate recall for non-ReAct mode
|
||||
if self.config.dataset_name in ("hotpotqa", "2wikimultihopqa"):
|
||||
recall_2, recall_5 = qa_judge.recall(sorted_context_ids, gold_file_ids)
|
||||
elif self.config.dataset_name == "musique":
|
||||
recall_2, recall_5 = qa_judge.recall(sorted_context, gold_file_ids)
|
||||
|
||||
result[f"{retriever.__class__.__name__ }_recall@2"] = recall_2
|
||||
result[f"{retriever.__class__.__name__ }_recall@5"] = recall_5
|
||||
|
||||
result_list.append(result)
|
||||
|
||||
self.save_results(result_list, [retriever.__class__.__name__ for retriever in retrievers])
|
||||
|
||||
|
||||
def save_results(self, result_list, retriever_names:List[str]):
|
||||
current_time = datetime.now()
|
||||
formatted_time = current_time.strftime("%Y%m%d%H%M%S")
|
||||
|
||||
dataset_name = self.config.dataset_name
|
||||
include_events = self.config.include_events
|
||||
include_concept = self.config.include_concept
|
||||
encoder_model_name = self.config.encoder_model_name
|
||||
reader_model_name = self.config.reader_model_name
|
||||
|
||||
# use last part of model name as identifier
|
||||
if "/" in encoder_model_name:
|
||||
encoder_model_name = encoder_model_name.split("/")[-1]
|
||||
if "/" in reader_model_name:
|
||||
reader_model_name = reader_model_name.split("/")[-1]
|
||||
|
||||
summary_file = f"./result/{dataset_name}/summary_{formatted_time}_event{include_events}_concept{include_concept}_{encoder_model_name}_{reader_model_name}.json"
|
||||
if not os.path.exists(os.path.dirname(summary_file)):
|
||||
os.makedirs(os.path.dirname(summary_file), exist_ok=True)
|
||||
|
||||
result_dir = f"./result/{dataset_name}/result_{formatted_time}_event{include_events}_concept{include_concept}_{encoder_model_name}_{reader_model_name}.json"
|
||||
if not os.path.exists(os.path.dirname(result_dir)):
|
||||
os.makedirs(os.path.dirname(result_dir), exist_ok=True)
|
||||
|
||||
summary_dict = self.calculate_summary(result_list, retriever_names)
|
||||
|
||||
with open(summary_file, "w") as f_summary:
|
||||
json.dump(summary_dict, f_summary)
|
||||
f_summary.write("\n")
|
||||
|
||||
with open(result_dir, "w") as f:
|
||||
for result in result_list:
|
||||
json.dump(result, f)
|
||||
f.write("\n")
|
||||
|
||||
def calculate_summary(self, result_list, method):
|
||||
summary_dict = {}
|
||||
for retriever_name in method:
|
||||
if not all(f"{retriever_name}_em" in result for result in result_list):
|
||||
raise ValueError(f"Missing {retriever_name}_em in results")
|
||||
if not all(f"{retriever_name}_f1" in result for result in result_list):
|
||||
raise ValueError(f"Missing {retriever_name}_f1 in results")
|
||||
|
||||
average_em = sum([result[f"{retriever_name}_em"] for result in result_list]) / len(result_list)
|
||||
average_f1 = sum([result[f"{retriever_name}_f1"] for result in result_list]) / len(result_list)
|
||||
|
||||
# Only calculate recall metrics if they exist in the results
|
||||
if all(f"{retriever_name}_recall@2" in result for result in result_list):
|
||||
average_recall_2 = sum([result[f"{retriever_name}_recall@2"] for result in result_list]) / len(result_list)
|
||||
average_recall_5 = sum([result[f"{retriever_name}_recall@5"] for result in result_list]) / len(result_list)
|
||||
summary_dict.update({
|
||||
f"{retriever_name}_average_f1": average_f1,
|
||||
f"{retriever_name}_average_em": average_em,
|
||||
f"{retriever_name}_average_recall@2": average_recall_2,
|
||||
f"{retriever_name}_average_recall@5": average_recall_5,
|
||||
})
|
||||
else:
|
||||
# For ReAct mode where recall metrics don't exist
|
||||
summary_dict.update({
|
||||
f"{retriever_name}_average_f1": average_f1,
|
||||
f"{retriever_name}_average_em": average_em,
|
||||
})
|
||||
|
||||
return summary_dict
|
||||
158
AIEC-RAG/atlas_rag/evaluation/evaluation.py
Normal file
158
AIEC-RAG/atlas_rag/evaluation/evaluation.py
Normal file
@ -0,0 +1,158 @@
|
||||
import re
|
||||
from collections import Counter
|
||||
from typing import Tuple
|
||||
import argparse
|
||||
import json
|
||||
from tqdm import tqdm
|
||||
class QAJudger:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def split_answer(self, generated_text):
|
||||
if "Answer:" in generated_text:
|
||||
generated_text = generated_text.split("Answer:")[-1]
|
||||
elif "answer:" in generated_text:
|
||||
generated_text = generated_text.split("answer:")[-1]
|
||||
# if answer is none
|
||||
if not generated_text:
|
||||
return "none"
|
||||
return generated_text
|
||||
|
||||
def normalize_answer(self, answer: str) -> str:
|
||||
"""Direct copy of the normalization from QAExactMatch/QAF1Score"""
|
||||
# Lowercase and normalize whitespace
|
||||
answer = answer.lower()
|
||||
# Replace hyphens with spaces
|
||||
answer = answer.replace('-', ' ')
|
||||
# Remove all other punctuation
|
||||
answer = re.sub(r'[^\w\s]', '', answer)
|
||||
# Standardize whitespace
|
||||
return ' '.join(answer.split())
|
||||
|
||||
def judge(self, generated_text: str, reference_text: str) -> Tuple[int, float]:
|
||||
"""Direct port of the original scoring logic"""
|
||||
# Extract answer from generated text
|
||||
pred_answer = self.split_answer(generated_text)
|
||||
|
||||
# Normalize both answers
|
||||
pred_norm = self.normalize_answer(pred_answer)
|
||||
ref_norm = self.normalize_answer(reference_text)
|
||||
|
||||
# Exact match calculation
|
||||
em = 1 if pred_norm == ref_norm else 0
|
||||
|
||||
# F1 calculation (direct port from QAF1Score)
|
||||
pred_tokens = pred_norm.split()
|
||||
ref_tokens = ref_norm.split()
|
||||
|
||||
common = Counter(pred_tokens) & Counter(ref_tokens)
|
||||
num_same = sum(common.values())
|
||||
|
||||
if num_same == 0:
|
||||
return em, 0.0
|
||||
|
||||
precision = num_same / len(pred_tokens) if pred_tokens else 0.0
|
||||
recall = num_same / len(ref_tokens) if ref_tokens else 0.0
|
||||
|
||||
if (precision + recall) == 0:
|
||||
f1 = 0.0
|
||||
else:
|
||||
f1 = 2 * (precision * recall) / (precision + recall)
|
||||
|
||||
return em, f1
|
||||
|
||||
def recall_at_k(self, retrieved_text: list, reference_text: list, k: int) -> float:
|
||||
"""Calculates recall at k based on the top k retrieved texts."""
|
||||
successful_retrievals = 0
|
||||
|
||||
# Limit the retrieved texts to the top k entries
|
||||
limited_retrieved_text = retrieved_text[:k]
|
||||
|
||||
for ref_text in reference_text:
|
||||
for ret_text in limited_retrieved_text:
|
||||
if ref_text in ret_text:
|
||||
successful_retrievals += 1
|
||||
break
|
||||
|
||||
recall = successful_retrievals / len(reference_text) if reference_text else 0
|
||||
return recall
|
||||
|
||||
# recall for 1 answer
|
||||
def recall(self, retrieved_text: list, reference_text: list) -> dict:
|
||||
"""Calculates recall values at different k levels."""
|
||||
recall_values = {
|
||||
'recall@2': self.recall_at_k(retrieved_text, reference_text, 2),
|
||||
'recall@5': self.recall_at_k(retrieved_text, reference_text, 5),
|
||||
}
|
||||
return recall_values['recall@2'], recall_values['recall@5']
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
argument_parser = argparse.ArgumentParser()
|
||||
argument_parser.add_argument("--file_path", type=str, required=True, help="Path to the JSON file containing results.")
|
||||
args = argument_parser.parse_args()
|
||||
|
||||
# Initialize the QAJudger
|
||||
llm_judge = QAJudger()
|
||||
|
||||
# Load results from the JSON file
|
||||
result_list = []
|
||||
with open(args.file_path, 'r') as file:
|
||||
for line in file:
|
||||
if line.strip(): # Make sure the line is not empty
|
||||
try:
|
||||
result = json.loads(line.strip())
|
||||
result_list.append(result)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"Error decoding JSON: {e}")
|
||||
|
||||
# Debugging output to inspect the loaded data structure
|
||||
# print("Loaded data structure:", result_list)
|
||||
|
||||
# Evaluate each entry in result_list
|
||||
for result in tqdm(result_list):
|
||||
if isinstance(result, dict): # Ensure each result is a dictionary
|
||||
question = result["question"]
|
||||
answer = result["answer"]
|
||||
|
||||
# Evaluate generated answers with Hippo and Hippo2
|
||||
hippo_generated_answer = result["hippo_generated_answer"]
|
||||
hippo2_generated_answer = result["hippo2_generated_answer"]
|
||||
|
||||
# Split and judge the answers
|
||||
hippo_short_answer = llm_judge.split_answer(hippo_generated_answer)
|
||||
hippo_em, hippo_f1 = llm_judge.judge(hippo_short_answer, answer)
|
||||
|
||||
hippo2_short_answer = llm_judge.split_answer(hippo2_generated_answer)
|
||||
hippo2_em, hippo2_f1 = llm_judge.judge(hippo2_short_answer, answer)
|
||||
|
||||
# Store the scores back in the result dictionary
|
||||
result["hippo_em"] = hippo_em
|
||||
result["hippo_f1"] = hippo_f1
|
||||
result["hippo2_em"] = hippo2_em
|
||||
result["hippo2_f1"] = hippo2_f1
|
||||
|
||||
result['recall@2'], result['recall@5'] = llm_judge.recall(result['hippo2_id'], result['gold_file_ids'])
|
||||
result['recall@2_hippo'], result['recall@5_hippo'] = llm_judge.recall(result['hippo_id'], result['gold_file_ids'])
|
||||
|
||||
# Calculate averages
|
||||
average_em_with_hippo = sum(result["hippo_em"] for result in result_list) / len(result_list)
|
||||
average_em_with_hippo2 = sum(result["hippo2_em"] for result in result_list) / len(result_list)
|
||||
|
||||
average_f1_with_hippo = sum(result["hippo_f1"] for result in result_list) / len(result_list)
|
||||
average_f1_with_hippo2 = sum(result["hippo2_f1"] for result in result_list) / len(result_list)
|
||||
|
||||
average_recall2_with_hippo = sum(result['recall@2'] for result in result_list) / len(result_list)
|
||||
average_recall5_with_hippo = sum(result['recall@5'] for result in result_list) / len(result_list)
|
||||
average_recall2 = sum(result['recall@2_hippo'] for result in result_list) / len(result_list)
|
||||
average_recall5 = sum(result['recall@5_hippo'] for result in result_list) / len(result_list)
|
||||
# Output the averages
|
||||
print(f"Average EM with Hippo: {average_em_with_hippo:.4f}")
|
||||
print(f"Average EM with Hippo2: {average_em_with_hippo2:.4f}")
|
||||
print(f"Average F1 with Hippo: {average_f1_with_hippo:.4f}")
|
||||
print(f"Average F1 with Hippo2: {average_f1_with_hippo2:.4f}")
|
||||
|
||||
print(f"Average Recall@2: {average_recall2:.4f}")
|
||||
print(f"Average Recall@5: {average_recall5:.4f}")
|
||||
print(f"Average Recall@2 with Hippo: {average_recall2_with_hippo:.4f}")
|
||||
print(f"Average Recall@5 with Hippo: {average_recall5_with_hippo:.4f}")
|
||||
Reference in New Issue
Block a user