Alignment-Lab-AI/StampyAI-alignment-data

1import json2from pathlib import Path3 4import datasets5from datasets import Value, Sequence, Features6 7 8_CITATION = '''9@article{kirchner2022understanding,10  title={Understanding AI Alignment Research: A Systematic Analysis},11  author={Kirchner, Jan H and Smith, Logan and Thibodeau, Jacques and McDonnell, Kyle and Reynolds, Laria},12  journal={arXiv preprint arXiv:2022.4338861},13  year={2022}14}15'''16 17_DESCRIPTION = """The AI Alignment Research Dataset is a collection of documents related to AI Alignment and Safety from various books, research papers, and alignment related blog posts."""18 19_HOMEPAGE = "https://github.com/StampyAI/alignment-research-dataset"20 21_LICENSE = "MIT license"22 23_VERSION_ = '0.0.0'24 25 26def iterate_file(filename):27    print(filename)28    with open(filename) as f:29        for l in f:30            try:31                yield json.loads(l)32            except Exception as e:33                print(f'Could not parse: {l}')34 35 36## Feature extractor helpers37def get_type(value):38    """Recursively get the huggingface type for the provided value."""39    if value is None:40        return None41    if value and isinstance(value, (tuple, list)):42        return features.Sequence(43            get_type(value[0])44        )45    if value and isinstance(value, dict):46        return {k: get_type(v) for k, v in value.items()}47    if isinstance(value, str):48        return Value('string')49    if isinstance(value, int):50        return Value('int32')51    if isinstance(value, float):52        return Value('double')53    if isinstance(value, bool):54        return Value('bool')55    return None56 57 58def print_extra_features(files):59    """Go through all the provided files, and get the non default features for the given file.60 61    This can be done manually but would be a hassle.62    It's assumed that the files contain a json object on each line.63    """64    ignored_keys = [65        'comments',  # Comments are arbitrarily nested objects, which doesn't play nice with huggingface66    ]67 68    per_file = {}69    for filename in sorted(files):70        extra_types = {}71        for item in iterate_file(filename):72            for k, v in item.items():73                if (k not in extra_types or not extra_types[k]) and k not in ignored_keys and k not in DEFAULT_FEATURES:74                    extra_types[k] = get_type(v)75        per_file[filename] = extra_types76 77    print('DATASOURCES = {')78    for k, features in per_file.items():79        vals = ',\n'.join(f"        '{k}': {v}" for k, v in features.items())80        print(f"    '{k.stem}': #\n{vals}\n    $,".replace('#', '{').replace('$', '}'))81    print('}')82 83 84# These keys are present in all files85DEFAULT_FEATURES = {86    'id': Value('string'),87    'source': Value('string'),88    'title': Value('string'),89    'text': Value('large_string'),90    'url': Value('string'),91    'date_published': Value(dtype='string'),92    'authors': Sequence(feature=Value(dtype='string'), length=-1),93    'summary': Sequence(feature=Value(dtype='string'), length=-1),94    'source_type': Value(dtype='string'),95}96 97 98# Per datasource additional features99DATASOURCES = {100    'agentmodels': {101        'book_title': Value(dtype='string'),102    },103    'agisf': {},104    'aisafety.info': {},105    'alignmentforum': {106        'karma': Value(dtype='int32'),107        'votes': Value(dtype='int32'),108        'words': Value(dtype='int32'),109        'comment_count': Value(dtype='int32'),110        'tags': Sequence(feature=Value(dtype='string')),111        'modified_at': Value(dtype='string'),112    },113    'arbital': {114        'alias': Value(dtype='string'),115        'tags': Sequence(feature=Value(dtype='string')),116    },117    'arxiv': {118        'data_last_modified': Value(dtype='string'),119        'abstract': Value(dtype='string'),120        'author_comment': Value(dtype='string'),121        'journal_ref': Value(dtype='string'),122        'doi': Value(dtype='string'),123        'primary_category': Value(dtype='string'),124        'categories': Sequence(feature=Value(dtype='string'), length=-1),125    },126    'blogs': {127        'initial_source': Value(dtype='string'),128    },129    'distill': {130        'abstract': Value(dtype='string'),131        'journal_ref': Value(dtype='string'),132        'doi': Value(dtype='string'),133        'bibliography_bib': Sequence(feature={'title': Value(dtype='string')}, length=-1),134    },135    'eaforum': {136        'karma': Value(dtype='int32'),137        'votes': Value(dtype='int32'),138        'words': Value(dtype='int32'),139        'comment_count': Value(dtype='int32'),140        'tags': Sequence(feature=Value(dtype='string')),141        'modified_at': Value(dtype='string'),142    },143    'lesswrong': {144        'karma': Value(dtype='int32'),145        'votes': Value(dtype='int32'),146        'words': Value(dtype='int32'),147        'comment_count': Value(dtype='int32'),148        'tags': Sequence(feature=Value(dtype='string')),149        'modified_at': Value(dtype='string'),150    },151    'special_docs': {},152    'youtube': {},153}154 155 156def join_features(features, to_join):157    """Recursively join the provided dicts.158 159    `to_join` can either be a dict to be merged, or a list of dicts to merge.160    """161    if not to_join:162        return Features(features)163    if isinstance(to_join, dict):164        return Features(dict(features, **to_join))165    return join_features(dict(features, **to_join[0]), to_join[1:])166 167 168class AlignmentResearchDatasetConfig(datasets.BuilderConfig):169    """BuilderConfig for AlignmentResaerchDataset."""170 171    def __init__(self, sources, features, **kwargs):172        """BuilderConfig for AlignmentResaerchDataset.173 174        :param List[string] sources: the sources which will be used by this config175        """176        super().__init__(version=datasets.Version(_VERSION_), **kwargs)177        self.sources = sources178        self.features = join_features(DEFAULT_FEATURES, features)179 180    @property181    def files(self):182        return [f'{source}.jsonl' for source in self.sources]183 184 185class AlignmentResaerchDataset(datasets.GeneratorBasedBuilder):186    VERSION = datasets.Version(_VERSION_)187 188    BUILDER_CONFIGS = [189        AlignmentResearchDatasetConfig(190            name='all',191            description='All data files',192            sources=list(DATASOURCES.keys()),193            features=list(DATASOURCES.values())194        )195    ] + [196        AlignmentResearchDatasetConfig(name=source, sources=[source], features=features) for source, features in DATASOURCES.items()197    ]198    DEFAULT_CONFIG_NAME = 'all'199 200    def _info(self):201        return datasets.DatasetInfo(202            description=_DESCRIPTION,203            features=self.config.features,204            homepage=_HOMEPAGE,205            license=_LICENSE,206            citation=_CITATION,207        )208 209    def _split_generators(self, dl_manager):210        downloaded_files = dl_manager.download_and_extract(self.config.files)211        return [212            datasets.SplitGenerator(213                name=datasets.Split.TRAIN,214                gen_kwargs={'files': downloaded_files}215            )216        ]217 218    # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`219    def _generate_examples(self, files):220        seen = set()221 222        def is_good(item):223            item_id = item and item.get('id')224            if not item_id or item_id in seen:225                return False226            seen.add(item_id)227 228            return item['text'] not in [None, '', 'n/a']229 230        def prepare_example(item):231            return item['id'], {k: item.get(k) for k in self.config.features}232 233        lines = (item for filename in files for item in iterate_file(filename))234        for item in map(prepare_example, filter(is_good, lines)):235            yield item236