""" Build the evaluation corpus and raw eval spec from HotpotQA + SQuAD v2. Outputs: - data/eval_corpus/*.txt one file per unique paragraph title - evaluation/data/eval_raw.json 100 unannotated eval cases (no chunk ids yet) Run download_eval_sources.py first. """ import hashlib import json import random import re from pathlib import Path import pandas as pd # ── constants ──────────────────────────────────────────────────────────────── N_HOTPOT = 80 N_FALLBACK = 20 RANDOM_SEED = 13 # Anchored paths — work from any working directory when run as a module _PACKAGE_DIR = Path(__file__).resolve().parent # evaluation/ _PROJECT_ROOT = _PACKAGE_DIR.parent # project root EVAL_SOURCES_DIR = _PROJECT_ROOT / "data" / "eval_sources" HOTPOT_PATH = EVAL_SOURCES_DIR / "hotpot_validation.parquet" SQUAD_PATH = EVAL_SOURCES_DIR / "squad_v2_dev.json" CORPUS_DIR = _PROJECT_ROOT / "data" / "eval_corpus" RAW_EVAL_PATH = _PACKAGE_DIR / "data" / "eval_raw.json" # ── helpers ────────────────────────────────────────────────────────────────── def slug(text: str) -> str: """Convert a title to a filesystem-safe slug (lowercase, hyphens).""" text = text.lower() text = re.sub(r"[^\w\s-]", "", text) text = re.sub(r"[\s_]+", "-", text.strip()) return text def title_hash(title: str) -> str: """Return the first 8 hex chars of the sha256 of a title string.""" return hashlib.sha256(title.encode("utf-8")).hexdigest()[:8] def corpus_filename(title: str) -> str: """Return a safe, unique filename for a paragraph with this title.""" return f"{slug(title)[:60]}-{title_hash(title)}.txt" def inspect_hotpot_schema(df: pd.DataFrame) -> None: """Print the columns and one row to verify the parquet schema.""" print("HotpotQA columns:", list(df.columns)) first = df.iloc[0] print(" question:", first["question"]) print(" answer:", first["answer"]) print(" type:", first["type"]) context = first["context"] print(" context type:", type(context)) print(" context keys:", list(context.keys()) if hasattr(context, "keys") else "N/A") titles = context["title"] sentences = context["sentences"] print(f" context paragraphs: {len(titles)}") print(f" first title: {titles[0]}") print(f" first para sentences[0]: {sentences[0][0][:80]!r}") sf = first["supporting_facts"] print(" supporting_facts keys:", list(sf.keys()) if hasattr(sf, "keys") else "N/A") print(" supporting_facts titles:", sf["title"]) print(" supporting_facts sent_id:", sf["sent_id"]) # ── step 1: load and sample hotpot ─────────────────────────────────────────── def load_hotpot_sample() -> list: """Load the HotpotQA parquet and return N_HOTPOT deterministically sampled rows.""" print(f"Loading {HOTPOT_PATH} ...") df = pd.read_parquet(HOTPOT_PATH) print(f" Loaded {len(df)} rows.") print("\n--- HotpotQA schema inspection ---") inspect_hotpot_schema(df) print("----------------------------------\n") rng = random.Random(RANDOM_SEED) indices = list(range(len(df))) rng.shuffle(indices) selected = [df.iloc[i] for i in indices[:N_HOTPOT]] return selected # ── step 2: collect unique paragraphs ──────────────────────────────────────── def collect_paragraphs(rows: list) -> dict[str, str]: """ Return {title: paragraph_text} for every unique title across all rows. If the same title appears in multiple questions, its sentence list from the first occurrence is used (they are identical in the dataset). """ paragraphs: dict[str, str] = {} for row in rows: context = row["context"] titles = context["title"] sentences_list = context["sentences"] for title, sentences in zip(titles, sentences_list): if title not in paragraphs: paragraphs[title] = " ".join(sentences) return paragraphs # ── step 3: write corpus files ─────────────────────────────────────────────── def write_corpus(paragraphs: dict[str, str]) -> None: """Clear data/eval_corpus/ and write one .txt file per paragraph.""" if CORPUS_DIR.exists(): for f in CORPUS_DIR.iterdir(): if f.suffix == ".txt": f.unlink() CORPUS_DIR.mkdir(parents=True, exist_ok=True) for title, content in paragraphs.items(): filename = corpus_filename(title) (CORPUS_DIR / filename).write_text(content, encoding="utf-8") print(f"Wrote {len(paragraphs)} corpus files to {CORPUS_DIR}/") # ── step 4: build hotpot eval records ──────────────────────────────────────── def build_hotpot_records(rows: list) -> tuple[list[dict], int]: """ Build one eval record per HotpotQA row. Returns (records, skipped_supporting_fact_count). """ records = [] skipped_facts = 0 for row in rows: context = row["context"] titles = context["title"] sentences_list = context["sentences"] sf = row["supporting_facts"] sf_titles = sf["title"] sf_sent_ids = sf["sent_id"] # Build a title -> sentence list index for this row. title_to_sentences: dict[str, list[str]] = {} for title, sentences in zip(titles, sentences_list): title_to_sentences[title] = list(sentences) supporting = [] for sf_title, sf_sent_id in zip(sf_titles, sf_sent_ids): sentence_list = title_to_sentences.get(sf_title, []) if sf_sent_id >= len(sentence_list): skipped_facts += 1 continue sentence_text = sentence_list[sf_sent_id] supporting.append({"title": sf_title, "sentence": sentence_text}) records.append({ "question": str(row["question"]), "answer": str(row["answer"]), "expected_behavior": "answer", "supporting": supporting, "source": "hotpotqa", "type": str(row["type"]), }) return records, skipped_facts # ── step 5: load squad unanswerable questions ───────────────────────────────── def load_squad_fallback_records(corpus_titles: set[str]) -> list[dict]: """ Return N_FALLBACK unanswerable SQuAD v2 questions that don't overlap corpus titles. Skips any question whose text contains a corpus title (case-insensitive). """ print(f"Loading {SQUAD_PATH} ...") squad_data = json.loads(SQUAD_PATH.read_text(encoding="utf-8")) corpus_titles_lower = {t.lower() for t in corpus_titles} impossible_questions: list[str] = [] for article in squad_data["data"]: for paragraph in article["paragraphs"]: for qa in paragraph["qas"]: if not qa.get("is_impossible", False): continue question_text = qa["question"] question_lower = question_text.lower() # Skip if any corpus title appears in the question text. if any(title in question_lower for title in corpus_titles_lower): continue impossible_questions.append(question_text) print(f" Found {len(impossible_questions)} unanswerable questions after title filter.") rng = random.Random(RANDOM_SEED) rng.shuffle(impossible_questions) selected = impossible_questions[:N_FALLBACK] records = [] for question in selected: records.append({ "question": question, "answer": "", "expected_behavior": "fallback", "supporting": [], "source": "squad_v2", }) return records # ── main ───────────────────────────────────────────────────────────────────── def main() -> None: random.seed(RANDOM_SEED) RAW_EVAL_PATH.parent.mkdir(parents=True, exist_ok=True) # Step 1: load hotpot sample hotpot_rows = load_hotpot_sample() # Step 2: collect unique paragraphs paragraphs = collect_paragraphs(hotpot_rows) # Step 3: write corpus files write_corpus(paragraphs) # Step 4: build hotpot eval records hotpot_records, skipped_facts = build_hotpot_records(hotpot_rows) print(f"Built {len(hotpot_records)} HotpotQA eval records.") if skipped_facts: print(f" Skipped {skipped_facts} supporting facts (sent_id out of range).") # Step 5: load squad fallback records squad_records = load_squad_fallback_records(set(paragraphs.keys())) # Step 6: combine and write combined = hotpot_records + squad_records RAW_EVAL_PATH.write_text( json.dumps(combined, indent=2, ensure_ascii=False), encoding="utf-8", ) print(f"\nWrote {len(combined)} records to {RAW_EVAL_PATH}") # Step 7: summary answer_count = sum(1 for r in combined if r["expected_behavior"] == "answer") fallback_count = sum(1 for r in combined if r["expected_behavior"] == "fallback") print("\n--- Summary ---") print(f" Corpus files: {len(paragraphs)}") print(f" Answer cases: {answer_count}") print(f" Fallback cases: {fallback_count}") print(f" Total records: {len(combined)}") if skipped_facts: print(f" Skipped supporting facts: {skipped_facts}") print("\nSample records (first 2):") for record in combined[:2]: print(json.dumps(record, indent=2, ensure_ascii=False)[:400]) print(" ...") if __name__ == "__main__": main()