#!/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()