minimind/scripts/validate_search_shortqa_dat...

79 lines
2.4 KiB
Python

import argparse
import json
from collections import Counter
def iter_jsonl(path):
with open(path, "r", encoding="utf-8") as f:
for i, line in enumerate(f, 1):
line = line.strip()
if line:
yield i, json.loads(line)
def as_list(value):
if value is None:
return []
if isinstance(value, list):
return value
return [value]
def validate_row(row):
errors = []
if not str(row.get("id", "")).strip():
errors.append("missing_id")
if not str(row.get("question", "")).strip():
errors.append("missing_question")
answers = row.get("answers") or row.get("answer")
if not as_list(answers):
errors.append("missing_answer")
contexts = row.get("contexts") or []
if not contexts:
errors.append("missing_contexts")
doc_ids = {str(doc.get("id", "")) for doc in contexts if isinstance(doc, dict)}
gold_ids = {str(x) for x in as_list(row.get("gold_doc_ids"))}
if not gold_ids:
errors.append("missing_gold_doc_ids")
elif not gold_ids.issubset(doc_ids):
errors.append("gold_doc_not_in_contexts")
empty_contexts = sum(1 for doc in contexts if not str(doc.get("text", "")).strip())
if empty_contexts:
errors.append("empty_context_text")
return errors, len(contexts), len(doc_ids), len(gold_ids)
def main():
parser = argparse.ArgumentParser(description="Validate SearchShortQA JSONL")
parser.add_argument("--data", required=True)
parser.add_argument("--max_errors", type=int, default=20)
args = parser.parse_args()
error_counts = Counter()
context_counts = Counter()
n = 0
examples = []
for lineno, row in iter_jsonl(args.data):
n += 1
errors, n_ctx, _, _ = validate_row(row)
context_counts[n_ctx] += 1
for err in errors:
error_counts[err] += 1
if errors and len(examples) < args.max_errors:
examples.append({"line": lineno, "id": row.get("id"), "errors": errors})
print(f"rows: {n}")
print("context_count_distribution:", dict(sorted(context_counts.items())))
print("error_counts:", dict(error_counts))
if examples:
print("error_examples:")
for ex in examples:
print(json.dumps(ex, ensure_ascii=False))
if error_counts:
raise SystemExit(1)
print("SearchShortQA validation passed.")
if __name__ == "__main__":
main()