#!/usr/bin/env python
"""Download the agreed multimodal test datasets into a local folder."""
from __future__ import annotations
import argparse
import json
import os
import shutil
import tarfile
import urllib.request
import urllib.parse
from pathlib import Path
from typing import Iterable, Optional
WIKITEXT_NAME = "wikitext"
WIKITEXT_CONFIG = "wikitext-103-v1"
CORD_DATASET = "naver-clova-ix/cord-v2"
LIBRISPEECH_URLS = [
"https://www.openslr.org/resources/12/dev-clean.tar.gz",
"https://www.openslr.org/resources/12/test-clean.tar.gz",
]
DEFAULT_VIDEO_URLS = [
"https://upload.wikimedia.org/wikipedia/commons/d/d2/The_Magic_of_Museums_-_Dynamic_Centers_of_Education_and_Innovation.webm",
"https://upload.wikimedia.org/wikipedia/commons/0/0b/IMD_2024Fascinating_World_of_Museum_Education.webm",
"https://upload.wikimedia.org/wikipedia/commons/2/26/IMD_2024_Museums_Important_Research_Institutions.webm",
"https://upload.wikimedia.org/wikipedia/commons/4/4f/Presentaci%C3%B3n_de_la_metodolog%C3%ADa_Great_Little_People.webm",
]
def ensure_dir(path: Path) -> None:
path.mkdir(parents=True, exist_ok=True)
def download_file(url: str, dest: Path) -> Path:
ensure_dir(dest.parent)
if dest.exists():
return dest
print(f"Downloading {url} -> {dest}")
request = urllib.request.Request(
url,
headers={"User-Agent": "CognitiveRAG/0.1 (dataset downloader)"},
)
with urllib.request.urlopen(request) as response, open(dest, "wb") as handle:
shutil.copyfileobj(response, handle)
return dest
def extract_tar_gz(archive_path: Path, dest_dir: Path) -> None:
print(f"Extracting {archive_path} -> {dest_dir}")
ensure_dir(dest_dir)
with tarfile.open(archive_path, "r:gz") as tar:
tar.extractall(dest_dir)
def download_librispeech(out_dir: Path) -> None:
ensure_dir(out_dir)
for url in LIBRISPEECH_URLS:
archive_path = out_dir / Path(url).name
download_file(url, archive_path)
extract_tar_gz(archive_path, out_dir)
def download_wikitext(out_dir: Path) -> None:
try:
from datasets import load_dataset
except Exception as exc:
raise RuntimeError("Missing datasets. Install with `pip install datasets`.") from exc
ensure_dir(out_dir)
ds = load_dataset(WIKITEXT_NAME, WIKITEXT_CONFIG)
for split in ds:
out_path = out_dir / f"{split}.txt"
with open(out_path, "w", encoding="utf-8") as handle:
for row in ds[split]:
text = (row.get("text") or "").rstrip()
if text:
handle.write(text + "\n")
print(f"Wrote WikiText-103 splits to {out_dir}")
def iter_dataset(ds, max_items: Optional[int] = None) -> Iterable[tuple[int, dict]]:
count = 0
for item in ds:
if max_items is not None and count >= max_items:
break
yield count, item
count += 1
def download_cord(out_dir: Path, max_items: Optional[int]) -> None:
try:
from datasets import load_dataset
except Exception as exc:
raise RuntimeError("Missing datasets. Install with `pip install datasets`.") from exc
ensure_dir(out_dir)
images_dir = out_dir / "images"
ensure_dir(images_dir)
metadata_path = out_dir / "metadata.jsonl"
ds = load_dataset(CORD_DATASET)
with open(metadata_path, "w", encoding="utf-8") as meta_handle:
for split_name, split in ds.items():
for idx, item in iter_dataset(split, max_items=max_items):
image = item.get("image")
if image is None:
continue
filename = f"{split_name}_{idx:06d}.png"
image_path = images_dir / filename
if not image_path.exists():
try:
image.save(image_path)
except Exception:
image.convert("RGB").save(image_path)
metadata = {k: v for k, v in item.items() if k != "image"}
metadata["image_file"] = filename
metadata["split"] = split_name
meta_handle.write(json.dumps(metadata, ensure_ascii=True, default=str) + "\n")
print(f"Wrote CORD-v2 images to {images_dir}")
def download_video_subset(out_dir: Path, list_path: Optional[Path]) -> None:
ensure_dir(out_dir)
if list_path is None:
urls = DEFAULT_VIDEO_URLS
else:
with open(list_path, "r", encoding="utf-8") as handle:
urls = [line.strip() for line in handle if line.strip()]
if not urls:
print("Video list is empty. Skipping video download.")
return
for url in urls:
url_path = urllib.parse.urlparse(url).path
if url_path.lower().endswith((".webm", ".mp4", ".ogv", ".mov", ".mkv")):
filename = Path(url_path).name
download_file(url, out_dir / filename)
continue
if shutil.which("yt-dlp") is None:
raise RuntimeError("Missing yt-dlp. Install with `pip install yt-dlp`.")
cmd = [
"yt-dlp",
"--no-playlist",
"-o",
str(out_dir / "%(id)s.%(ext)s"),
url,
]
print(f"Downloading video: {url}")
result = shutil.which("yt-dlp")
if result is None:
raise RuntimeError("yt-dlp not available.")
os.system(" ".join(cmd))
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Download multimodal test datasets.")
parser.add_argument(
"--out",
default="00_Data_test",
help="Output directory for downloaded datasets.",
)
parser.add_argument("--text", action="store_true", help="Download WikiText-103.")
parser.add_argument("--image", action="store_true", help="Download CORD-v2.")
parser.add_argument("--audio", action="store_true", help="Download LibriSpeech dev/test clean.")
parser.add_argument("--video", action="store_true", help="Download video subset via yt-dlp.")
parser.add_argument(
"--video-list",
default=None,
help="Path to newline-delimited video URLs/IDs (required if --video).",
)
parser.add_argument(
"--cord-max-items",
type=int,
default=None,
help="Optional cap on number of CORD examples per split.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
out_dir = Path(args.out)
selected = args.text or args.image or args.audio or args.video
if not selected:
args.text = args.image = args.audio = True
if args.text:
download_wikitext(out_dir / "text" / "wikitext-103")
if args.image:
download_cord(out_dir / "image" / "cord-v2", max_items=args.cord_max_items)
if args.audio:
download_librispeech(out_dir / "audio" / "librispeech")
if args.video:
list_path = Path(args.video_list) if args.video_list else None
download_video_subset(out_dir / "video" / "yt-ugc-subset", list_path)
if __name__ == "__main__":
main()