GotThatData/STARGATE

1import os2import json3import datasets4 5logger = datasets.logging.get_logger(__name__)6 7_CITATION = """\8@dataset{gotthatdata_stargate_2024,9  title  = {STARGATE: CIA Remote Viewing Archive},10  author = {GotThatData},11  year   = {2024},12  url    = {https://huggingface.co/datasets/GotThatData/STARGATE}13}14"""15 16_DESCRIPTION = """\17STARGATE is a dataset of 12,000+ declassified CIA PDFs related to remote viewing (RV), extrasensory perception (ESP), and anomalous cognition. 18This loader includes structured metadata and binary access to the original scanned PDFs.19"""20 21_HOMEPAGE = "https://huggingface.co/datasets/GotThatData/STARGATE"22 23_LICENSE = "CC-BY-4.0"24 25class StargatePDFConfig(datasets.BuilderConfig):26    def __init__(self, **kwargs):27        super(StargatePDFConfig, self).__init__(**kwargs)28 29class StargatePDFDataset(datasets.GeneratorBasedBuilder):30    VERSION = datasets.Version("1.0.0")31    BUILDER_CONFIGS = [32        StargatePDFConfig(name="default", version=VERSION, description="STARGATE raw PDFs with metadata")33    ]34 35    def _info(self):36        return datasets.DatasetInfo(37            description=_DESCRIPTION,38            features=datasets.Features({39                "filename": datasets.Value("string"),40                "document_id": datasets.Value("string"),41                "page_count": datasets.Value("int32"),42                "image_count": datasets.Value("int32"),43                "processed_at": datasets.Value("string"),44                "ocr_status": datasets.Value("string"),45                "text_extracted": datasets.Value("bool"),46                "source": datasets.Value("string"),47                "tags": datasets.Sequence(datasets.Value("string")),48                "pdf": datasets.Value("binary"),49            }),50            supervised_keys=None,51            homepage=_HOMEPAGE,52            license=_LICENSE,53            citation=_CITATION,54        )55 56    def _split_generators(self, dl_manager):57        # Path to local or downloaded files58        archive_path = dl_manager.download_and_extract("./")  # Change if remote59        metadata_path = os.path.join(archive_path, "metadata.json")60        data_dir = os.path.join(archive_path, "data")61 62        return [63            datasets.SplitGenerator(64                name=datasets.Split.TRAIN,65                gen_kwargs={"metadata_path": metadata_path, "data_dir": data_dir}66            )67        ]68 69    def _generate_examples(self, metadata_path, data_dir):70        logger.info(f"⏳ Loading metadata from {metadata_path}")71        with open(metadata_path, "r", encoding="utf-8") as f:72            records = json.load(f)73 74        for idx, record in enumerate(records):75            pdf_path = os.path.join(data_dir, record["filename"])76            if not os.path.isfile(pdf_path):77                logger.warning(f"🚫 Missing PDF: {pdf_path}")78                continue79 80            with open(pdf_path, "rb") as pdf_file:81                yield idx, {82                    "filename": record.get("filename"),83                    "document_id": record.get("document_id", record["filename"].replace(".pdf", "")),84                    "page_count": record.get("page_count", 0),85                    "image_count": record.get("image_count", 0),86                    "processed_at": record.get("processed_at", ""),87                    "ocr_status": record.get("ocr_status", "pending"),88                    "text_extracted": record.get("text_extracted", False),89                    "source": record.get("source", "CIA Stargate Archive"),90                    "tags": record.get("tags", []),91                    "pdf": pdf_file.read(),92                }93