Spaces:
Runtime error
Runtime error
File size: 3,806 Bytes
6dfa658 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | from __future__ import annotations
import argparse
import json
from pathlib import Path
CORPUS_FILES = ["real_corpus.jsonl", "corpus_index.jsonl", "corpus.jsonl"]
BENCHMARK_FILES = ["eval_qa.jsonl", "custom_benchmark.jsonl", "benchmark.jsonl", "eval_qa_150.jsonl"]
def read_jsonl(path: Path) -> list[dict]:
rows = []
with path.open("r", encoding="utf-8") as f:
for line_no, line in enumerate(f, start=1):
if not line.strip():
continue
try:
rows.append(json.loads(line))
except json.JSONDecodeError as exc:
raise SystemExit(f"{path}:{line_no} is not valid JSON: {exc}") from exc
return rows
def find_existing(data_dir: Path, filenames: list[str]) -> Path | None:
for filename in filenames:
path = data_dir / filename
if path.exists() and path.stat().st_size > 0:
return path
return None
def main() -> None:
parser = argparse.ArgumentParser(description="Validate custom corpus and benchmark files for the RAG system")
parser.add_argument("--data-dir", type=Path, required=True)
parser.add_argument("--require-benchmark", action="store_true")
args = parser.parse_args()
corpus_path = find_existing(args.data_dir, CORPUS_FILES)
if corpus_path is None:
raise SystemExit(f"No corpus file found in {args.data_dir}. Expected one of: {', '.join(CORPUS_FILES)}")
corpus = read_jsonl(corpus_path)
ids = set()
for idx, row in enumerate(corpus, start=1):
for key in ["id", "text"]:
if key not in row or not str(row[key]).strip():
raise SystemExit(f"{corpus_path}:{idx} missing required field '{key}'")
if row["id"] in ids:
raise SystemExit(f"{corpus_path}:{idx} duplicate id '{row['id']}'")
ids.add(row["id"])
benchmark_path = find_existing(args.data_dir, BENCHMARK_FILES)
benchmark_count = 0
missing_gold_ids = []
if benchmark_path:
benchmark = read_jsonl(benchmark_path)
benchmark_count = len(benchmark)
for idx, row in enumerate(benchmark, start=1):
question = row.get("question") or row.get("query")
if not question:
raise SystemExit(f"{benchmark_path}:{idx} missing 'question' or 'query'")
gold_ids = []
for key in ["source_id", "gold_source_id"]:
if row.get(key):
gold_ids.append(row[key])
for key in ["gold_chunk_ids", "relevant_documents", "relevant_doc_ids"]:
value = row.get(key)
if isinstance(value, str):
gold_ids.append(value)
elif isinstance(value, list):
gold_ids.extend(value)
for source in row.get("gold_sources", []) or []:
if source.get("corpus_row_id"):
gold_ids.append(source["corpus_row_id"])
for doc_id in gold_ids:
if doc_id not in ids:
missing_gold_ids.append((idx, doc_id))
elif args.require_benchmark:
raise SystemExit(f"No benchmark file found in {args.data_dir}. Expected one of: {', '.join(BENCHMARK_FILES)}")
if missing_gold_ids:
preview = ", ".join(f"row {row}: {doc_id}" for row, doc_id in missing_gold_ids[:10])
raise SystemExit(f"Benchmark references document ids not found in corpus: {preview}")
print("Custom data validation passed")
print(f"corpus_file: {corpus_path}")
print(f"corpus_documents: {len(corpus)}")
if benchmark_path:
print(f"benchmark_file: {benchmark_path}")
print(f"benchmark_questions: {benchmark_count}")
else:
print("benchmark_file: not provided")
if __name__ == "__main__":
main()
|