Spaces:
Sleeping
Sleeping
| import argparse | |
| import importlib.util | |
| import json | |
| import re | |
| import string | |
| from collections import Counter | |
| from pathlib import Path | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from config import CONFIG | |
| from inference import RAGEngine | |
| def _load_data_loader_module(): | |
| module_path = Path(__file__).with_name("data-loader.py") | |
| spec = importlib.util.spec_from_file_location("data_loader", module_path) | |
| if spec is None or spec.loader is None: | |
| raise RuntimeError("Could not load data-loader.py") | |
| module = importlib.util.module_from_spec(spec) | |
| spec.loader.exec_module(module) | |
| return module | |
| def normalize_text(text: str) -> str: | |
| text = text.lower() | |
| text = text.translate(str.maketrans("", "", string.punctuation)) | |
| text = re.sub(r"\b(a|an|the)\b", " ", text) | |
| text = re.sub(r"\s+", " ", text) | |
| return text.strip() | |
| def exact_match_score(prediction: str, ground_truth: str) -> float: | |
| return float(normalize_text(prediction) == normalize_text(ground_truth)) | |
| def f1_score(prediction: str, ground_truth: str) -> float: | |
| pred_tokens = normalize_text(prediction).split() | |
| truth_tokens = normalize_text(ground_truth).split() | |
| if not pred_tokens and not truth_tokens: | |
| return 1.0 | |
| if not pred_tokens or not truth_tokens: | |
| return 0.0 | |
| common = Counter(pred_tokens) & Counter(truth_tokens) | |
| num_same = sum(common.values()) | |
| if num_same == 0: | |
| return 0.0 | |
| precision = num_same / len(pred_tokens) | |
| recall = num_same / len(truth_tokens) | |
| return 2 * precision * recall / (precision + recall) | |
| def build_examples(limit: int) -> list[dict]: | |
| data_loader = _load_data_loader_module() | |
| from datasets import load_dataset | |
| dataset = load_dataset(CONFIG.dataset_name, split=CONFIG.dataset_question_split) | |
| examples = [] | |
| for row in dataset: | |
| answers = row.get("answers", {}) | |
| answer_texts = answers.get("text", []) if isinstance(answers, dict) else [] | |
| if not answer_texts: | |
| continue | |
| examples.append( | |
| { | |
| "question": str(row[CONFIG.dataset_question_field]), | |
| "answers": [str(answer) for answer in answer_texts], | |
| "context": str(row.get(CONFIG.dataset_context_field, "")), | |
| } | |
| ) | |
| if len(examples) >= limit: | |
| break | |
| if not examples: | |
| questions = data_loader.load_questions() | |
| for question in questions[:limit]: | |
| examples.append({"question": question, "answers": [question], "context": ""}) | |
| return examples | |
| def evaluate(limit: int, *, top_k: int | None = None) -> dict: | |
| engine = RAGEngine() | |
| examples = build_examples(limit) | |
| # Align retrieval corpus to evaluation contexts so the gold answers are reachable. | |
| eval_documents = [] | |
| seen_contexts = set() | |
| for example in examples: | |
| ctx = example.get("context", "").strip() | |
| if not ctx or ctx in seen_contexts: | |
| continue | |
| seen_contexts.add(ctx) | |
| eval_documents.append(ctx) | |
| if eval_documents: | |
| engine.documents = eval_documents | |
| engine.chunks = engine._chunk_documents(eval_documents) | |
| engine.vectorizer = TfidfVectorizer(stop_words="english") | |
| engine.chunk_tfidf_matrix = engine.vectorizer.fit_transform([chunk["text"] for chunk in engine.chunks]) | |
| engine.chunk_embeddings = engine._encode_texts([chunk["text"] for chunk in engine.chunks]) | |
| results = [] | |
| em_total = 0.0 | |
| f1_total = 0.0 | |
| for example in examples: | |
| prediction = engine.answer(example["question"], top_k=top_k) | |
| pred_text = prediction["answer"] | |
| answers = example["answers"] | |
| em = max(exact_match_score(pred_text, answer) for answer in answers) | |
| f1 = max(f1_score(pred_text, answer) for answer in answers) | |
| em_total += em | |
| f1_total += f1 | |
| results.append( | |
| { | |
| "question": example["question"], | |
| "prediction": pred_text, | |
| "gold_answers": answers, | |
| "gold_context": example.get("context", ""), | |
| "exact_match": em, | |
| "f1": f1, | |
| "answer_score": prediction["answer_score"], | |
| "context": prediction["context"], | |
| } | |
| ) | |
| count = len(results) or 1 | |
| report = { | |
| "dataset": CONFIG.dataset_name, | |
| "split": CONFIG.dataset_question_split, | |
| "examples": len(results), | |
| "exact_match": em_total / count, | |
| "f1": f1_total / count, | |
| "results": results, | |
| } | |
| return report | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Evaluate the RAG pipeline on a Hugging Face QA dataset.") | |
| parser.add_argument("--limit", type=int, default=25, help="Number of evaluation examples to score.") | |
| parser.add_argument("--top-k", type=int, default=None, help="Override the number of retrieved passages used during evaluation.") | |
| parser.add_argument("--fast", action="store_true", help="Run a smaller and faster evaluation using fewer examples and passages.") | |
| parser.add_argument( | |
| "--output", | |
| type=str, | |
| default=str(CONFIG.output_path.parent / "evaluation_report.json"), | |
| help="Path to write the evaluation report.", | |
| ) | |
| args = parser.parse_args() | |
| limit = args.limit | |
| top_k = args.top_k | |
| if args.fast: | |
| limit = min(limit, 2) | |
| top_k = 1 if top_k is None else min(top_k, 1) | |
| report = evaluate(limit, top_k=top_k) | |
| output_path = Path(args.output) | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| output_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8") | |
| print(f"Evaluated {report['examples']} examples") | |
| print(f"Exact Match: {report['exact_match']:.4f}") | |
| print(f"F1: {report['f1']:.4f}") | |
| print(f"Saved report to {output_path}") | |
| if __name__ == "__main__": | |
| main() | |