#!/usr/bin/env python3
"""
fetch-squad.py -- Download SQuAD dataset and extract Q&A as plain text sentences.
SQuAD (Stanford Question Answering Dataset) contains ~100K question-answer pairs
from Wikipedia paragraphs. We extract questions and context paragraphs as plain text.
Output: corpus/incoming/stage2-questions/*.txt (one file per topic chunk)
Usage: python3 mvp/scripts/fetch-squad.py
"""
import json
import os
import sys
import urllib.request
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
MVP_DIR = os.path.dirname(SCRIPT_DIR)
OUTPUT_DIR = os.path.join(MVP_DIR, "corpus", "incoming", "stage2-questions")
WORK_DIR = os.path.join(MVP_DIR, "corpus", ".work", "squad")
SQUAD_TRAIN_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v2.0.json"
SQUAD_DEV_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v2.0.json"
def download_file(url, dest):
"""Download a file if it doesn't already exist."""
if os.path.exists(dest):
print(f" Already downloaded: {dest}")
return
print(f" Downloading {url}...")
urllib.request.urlretrieve(url, dest)
size_mb = os.path.getsize(dest) / (1024 * 1024)
print(f" Downloaded: {size_mb:.1f} MB")
def extract_qa(data):
"""Extract questions and context paragraphs from SQuAD JSON.
Yields (title, question, context) tuples.
"""
for article in data["data"]:
title = article["title"]
for paragraph in article["paragraphs"]:
context = paragraph["context"]
for qa in paragraph["qas"]:
question = qa["question"]
yield (title, question, context)
def write_topic_files(qa_pairs, output_dir, pairs_per_file=500):
"""Write Q&A pairs to text files, grouped by count."""
os.makedirs(output_dir, exist_ok=True)
file_num = 0
pair_count = 0
total = 0
current_file = None
for title, question, context in qa_pairs:
if pair_count >= pairs_per_file or current_file is None:
if current_file:
current_file.close()
fname = os.path.join(output_dir, f"squad-{file_num:04d}.txt")
current_file = open(fname, "w", encoding="utf-8")
file_num += 1
pair_count = 0
# Write question as its own line
current_file.write(question.strip() + "\n")
# Write context sentences (split on period)
for sentence in context.split(". "):
sentence = sentence.strip().rstrip(".")
if len(sentence) > 20:
current_file.write(sentence + "\n")
current_file.write("\n") # Blank line between Q&A pairs
pair_count += 1
total += 1
if current_file:
current_file.close()
return file_num, total
def main():
print("═══ SQuAD Q&A Fetch Pipeline ═══")
print(f" Output dir: {OUTPUT_DIR}")
print(f" Work dir: {WORK_DIR}")
os.makedirs(WORK_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)
# Step 1: Download SQuAD datasets
train_file = os.path.join(WORK_DIR, "train-v2.0.json")
dev_file = os.path.join(WORK_DIR, "dev-v2.0.json")
download_file(SQUAD_TRAIN_URL, train_file)
download_file(SQUAD_DEV_URL, dev_file)
# Step 2: Extract Q&A pairs
print(" Extracting Q&A pairs...")
all_pairs = []
with open(train_file, "r", encoding="utf-8") as f:
train_data = json.load(f)
train_pairs = list(extract_qa(train_data))
print(f" Train: {len(train_pairs)} Q&A pairs")
all_pairs.extend(train_pairs)
with open(dev_file, "r", encoding="utf-8") as f:
dev_data = json.load(f)
dev_pairs = list(extract_qa(dev_data))
print(f" Dev: {len(dev_pairs)} Q&A pairs")
all_pairs.extend(dev_pairs)
# Step 3: Write to text files
print(" Writing text files...")
n_files, n_pairs = write_topic_files(all_pairs, OUTPUT_DIR)
print(f"\nDone! Files in: {OUTPUT_DIR}/")
print(f" {n_files} files")
print(f" {n_pairs} Q&A pairs")
print(f'\nTo process: (crawl-dir "{OUTPUT_DIR}")')
if __name__ == "__main__":
main()