Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import json | |
| import tempfile | |
| import urllib.request | |
| import zipfile | |
| from collections import defaultdict | |
| from collections.abc import Sequence | |
| from pathlib import Path | |
| from typing import Any | |
| from pydantic import BaseModel, ConfigDict, Field, field_validator | |
| XQUAD_VI_URL = "https://raw.githubusercontent.com/google-deepmind/xquad/master/xquad.vi.json" | |
| BEIR_SCIFACT_URL = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip" | |
| class ExternalCorpusDocument(BaseModel): | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| id: str = Field(min_length=1) | |
| text: str = Field(min_length=1) | |
| title: str | None = None | |
| def _strip_required(cls, value: str) -> str: | |
| stripped = value.strip() | |
| if not stripped: | |
| raise ValueError("Value cannot be blank.") | |
| return stripped | |
| def _strip_title(cls, value: str | None) -> str | None: | |
| if value is None: | |
| return None | |
| return value.strip() or None | |
| class ExternalBenchmarkQuery(BaseModel): | |
| model_config = ConfigDict(extra="forbid", frozen=True) | |
| id: str = Field(min_length=1) | |
| question: str = Field(min_length=1) | |
| ground_truth: str = "" | |
| gold_document_ids: list[str] = Field(min_length=1) | |
| answerable: bool = True | |
| def _strip_required(cls, value: str) -> str: | |
| stripped = value.strip() | |
| if not stripped: | |
| raise ValueError("Value cannot be blank.") | |
| return stripped | |
| def _strip_ground_truth(cls, value: str) -> str: | |
| return value.strip() | |
| def _normalize_gold_ids(cls, values: list[str]) -> list[str]: | |
| normalized = list(dict.fromkeys(value.strip() for value in values if value.strip())) | |
| if not normalized: | |
| raise ValueError("gold_document_ids cannot be empty.") | |
| return normalized | |
| def _write_jsonl(path: Path, records: Sequence[BaseModel]) -> None: | |
| path.write_text( | |
| "".join(json.dumps(record.model_dump(), ensure_ascii=False) + "\n" for record in records), | |
| encoding="utf-8", | |
| ) | |
| def write_external_benchmark( | |
| output_dir: Path, | |
| *, | |
| benchmark: str, | |
| source: str, | |
| corpus: list[ExternalCorpusDocument], | |
| queries: list[ExternalBenchmarkQuery], | |
| ) -> dict[str, object]: | |
| if not corpus: | |
| raise ValueError("External benchmark corpus cannot be empty.") | |
| if not queries: | |
| raise ValueError("External benchmark queries cannot be empty.") | |
| _validate_gold_documents(corpus, queries) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| _write_jsonl(output_dir / "corpus.jsonl", corpus) | |
| _write_jsonl(output_dir / "queries.jsonl", queries) | |
| manifest: dict[str, object] = { | |
| "schema_version": 1, | |
| "benchmark": benchmark, | |
| "source": source, | |
| "relevance_level": "document", | |
| "corpus_documents": len(corpus), | |
| "queries": len(queries), | |
| "answerable_queries": sum(query.answerable for query in queries), | |
| } | |
| (output_dir / "manifest.json").write_text( | |
| json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", | |
| encoding="utf-8", | |
| ) | |
| return manifest | |
| def _validate_gold_documents( | |
| corpus: list[ExternalCorpusDocument], | |
| queries: list[ExternalBenchmarkQuery], | |
| ) -> None: | |
| corpus_ids = {document.id for document in corpus} | |
| missing_gold = sorted( | |
| { | |
| document_id | |
| for query in queries | |
| for document_id in query.gold_document_ids | |
| if document_id not in corpus_ids | |
| } | |
| ) | |
| if missing_gold: | |
| preview = ", ".join(missing_gold[:5]) | |
| raise ValueError(f"Gold document ids are missing from the corpus: {preview}") | |
| def load_external_benchmark( | |
| benchmark_dir: Path, | |
| ) -> tuple[list[ExternalCorpusDocument], list[ExternalBenchmarkQuery], dict[str, Any]]: | |
| corpus_path = benchmark_dir / "corpus.jsonl" | |
| queries_path = benchmark_dir / "queries.jsonl" | |
| manifest_path = benchmark_dir / "manifest.json" | |
| for path in (corpus_path, queries_path, manifest_path): | |
| if not path.exists(): | |
| raise FileNotFoundError(f"External benchmark artifact not found: {path}") | |
| corpus = [ | |
| ExternalCorpusDocument.model_validate(json.loads(line)) | |
| for line in corpus_path.read_text(encoding="utf-8").splitlines() | |
| if line.strip() | |
| ] | |
| queries = [ | |
| ExternalBenchmarkQuery.model_validate(json.loads(line)) | |
| for line in queries_path.read_text(encoding="utf-8").splitlines() | |
| if line.strip() | |
| ] | |
| manifest_payload = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| if not isinstance(manifest_payload, dict): | |
| raise ValueError("External benchmark manifest must contain an object.") | |
| if not corpus or not queries: | |
| raise ValueError("External benchmark artifacts cannot be empty.") | |
| _validate_gold_documents(corpus, queries) | |
| return corpus, queries, manifest_payload | |
| def convert_xquad_vi_payload(payload: dict[str, Any], output_dir: Path) -> dict[str, object]: | |
| corpus: list[ExternalCorpusDocument] = [] | |
| queries: list[ExternalBenchmarkQuery] = [] | |
| articles = payload.get("data") | |
| if not isinstance(articles, list): | |
| raise ValueError("XQuAD payload must contain a data list.") | |
| for article_index, raw_article in enumerate(articles, start=1): | |
| if not isinstance(raw_article, dict): | |
| raise ValueError("XQuAD articles must be objects.") | |
| title = str(raw_article.get("title") or "").strip() or None | |
| paragraphs = raw_article.get("paragraphs") | |
| if not isinstance(paragraphs, list): | |
| raise ValueError("XQuAD article must contain a paragraphs list.") | |
| for paragraph_index, raw_paragraph in enumerate(paragraphs, start=1): | |
| if not isinstance(raw_paragraph, dict): | |
| raise ValueError("XQuAD paragraphs must be objects.") | |
| context = str(raw_paragraph.get("context") or "").strip() | |
| document_id = f"xquad-vi-{article_index:03d}-{paragraph_index:03d}" | |
| corpus.append(ExternalCorpusDocument(id=document_id, text=context, title=title)) | |
| qas = raw_paragraph.get("qas") | |
| if not isinstance(qas, list): | |
| raise ValueError("XQuAD paragraph must contain a qas list.") | |
| for raw_qa in qas: | |
| if not isinstance(raw_qa, dict): | |
| raise ValueError("XQuAD QA entries must be objects.") | |
| answers = raw_qa.get("answers") | |
| if not isinstance(answers, list) or not answers or not isinstance(answers[0], dict): | |
| raise ValueError("XQuAD QA entry must contain at least one answer.") | |
| queries.append( | |
| ExternalBenchmarkQuery( | |
| id=f"xquad-vi-{raw_qa.get('id')}", | |
| question=str(raw_qa.get("question") or ""), | |
| ground_truth=str(answers[0].get("text") or ""), | |
| gold_document_ids=[document_id], | |
| ) | |
| ) | |
| return write_external_benchmark( | |
| output_dir, | |
| benchmark="xquad-vi", | |
| source=XQUAD_VI_URL, | |
| corpus=corpus, | |
| queries=queries, | |
| ) | |
| def _read_jsonl_objects(path: Path) -> list[dict[str, Any]]: | |
| records: list[dict[str, Any]] = [] | |
| for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1): | |
| if not line.strip(): | |
| continue | |
| payload = json.loads(line) | |
| if not isinstance(payload, dict): | |
| raise ValueError(f"Expected object in {path} at line {line_number}.") | |
| records.append(payload) | |
| return records | |
| def convert_beir_scifact_directory(source_dir: Path, output_dir: Path) -> dict[str, object]: | |
| corpus_paths = list(source_dir.rglob("corpus.jsonl")) | |
| query_paths = list(source_dir.rglob("queries.jsonl")) | |
| qrels_paths = list(source_dir.rglob("qrels/test.tsv")) | |
| if len(corpus_paths) != 1 or len(query_paths) != 1 or len(qrels_paths) != 1: | |
| raise ValueError( | |
| "SciFact source must contain corpus.jsonl, queries.jsonl, and qrels/test.tsv." | |
| ) | |
| corpus = [ | |
| ExternalCorpusDocument( | |
| id=str(record.get("_id") or ""), | |
| title=str(record.get("title") or "") or None, | |
| text="\n\n".join( | |
| part | |
| for part in ( | |
| str(record.get("title") or "").strip(), | |
| str(record.get("text") or "").strip(), | |
| ) | |
| if part | |
| ), | |
| ) | |
| for record in _read_jsonl_objects(corpus_paths[0]) | |
| ] | |
| query_text = { | |
| str(record.get("_id") or ""): str(record.get("text") or "") | |
| for record in _read_jsonl_objects(query_paths[0]) | |
| } | |
| qrels: dict[str, list[str]] = defaultdict(list) | |
| with qrels_paths[0].open(encoding="utf-8", newline="") as input_file: | |
| for row in csv.DictReader(input_file, delimiter="\t"): | |
| if int(row["score"]) > 0: | |
| qrels[row["query-id"]].append(row["corpus-id"]) | |
| queries = [ | |
| ExternalBenchmarkQuery( | |
| id=f"beir-scifact-{query_id}", | |
| question=query_text[query_id], | |
| gold_document_ids=gold_ids, | |
| ) | |
| for query_id, gold_ids in sorted(qrels.items()) | |
| if query_id in query_text | |
| ] | |
| return write_external_benchmark( | |
| output_dir, | |
| benchmark="beir-scifact", | |
| source=BEIR_SCIFACT_URL, | |
| corpus=corpus, | |
| queries=queries, | |
| ) | |
| def _download(url: str, destination: Path) -> None: | |
| request = urllib.request.Request(url, headers={"User-Agent": "DocuLearn-RAG-evaluator/1.0"}) | |
| with urllib.request.urlopen(request, timeout=120) as response: # noqa: S310 | |
| destination.write_bytes(response.read()) | |
| def _safe_extract(archive: Path, destination: Path) -> None: | |
| destination_root = destination.resolve() | |
| with zipfile.ZipFile(archive) as zip_file: | |
| for member in zip_file.infolist(): | |
| member_path = (destination / member.filename).resolve() | |
| if not member_path.is_relative_to(destination_root): | |
| raise ValueError(f"Unsafe path in benchmark archive: {member.filename}") | |
| zip_file.extractall(destination) | |
| def prepare_xquad_vi(output_dir: Path, source: Path | None = None) -> dict[str, object]: | |
| if source is None: | |
| with tempfile.TemporaryDirectory(prefix="doculearn-xquad-") as temp_dir: | |
| source = Path(temp_dir) / "xquad.vi.json" | |
| _download(XQUAD_VI_URL, source) | |
| payload = json.loads(source.read_text(encoding="utf-8")) | |
| else: | |
| payload = json.loads(source.read_text(encoding="utf-8")) | |
| if not isinstance(payload, dict): | |
| raise ValueError("XQuAD source must contain a JSON object.") | |
| return convert_xquad_vi_payload(payload, output_dir) | |
| def prepare_beir_scifact(output_dir: Path, source: Path | None = None) -> dict[str, object]: | |
| with tempfile.TemporaryDirectory(prefix="doculearn-scifact-") as temp_dir: | |
| temp_path = Path(temp_dir) | |
| if source is None: | |
| archive = temp_path / "scifact.zip" | |
| _download(BEIR_SCIFACT_URL, archive) | |
| _safe_extract(archive, temp_path / "source") | |
| source_dir = temp_path / "source" | |
| elif source.is_file(): | |
| _safe_extract(source, temp_path / "source") | |
| source_dir = temp_path / "source" | |
| else: | |
| source_dir = source | |
| return convert_beir_scifact_directory(source_dir, output_dir) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Prepare a public retrieval benchmark.") | |
| parser.add_argument("benchmark", choices=("xquad-vi", "beir-scifact")) | |
| parser.add_argument("--output-dir", type=Path, required=True) | |
| parser.add_argument("--source", type=Path) | |
| args = parser.parse_args() | |
| if args.benchmark == "xquad-vi": | |
| manifest = prepare_xquad_vi(args.output_dir, args.source) | |
| else: | |
| manifest = prepare_beir_scifact(args.output_dir, args.source) | |
| print(json.dumps(manifest, ensure_ascii=False, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |