accessible_worlds.aw_publish_dataset

 1import argparse
 2from pathlib import Path
 3import random
 4import orjson
 5from datasets import Dataset, DatasetDict
 6
 7from .aw_utils import load_json
 8
 9def curate_row( r ):
10    return {
11        "text" : r["completion"],
12        "id"   : r["index"],
13        "explanitory_style": r["explanitory_style"],
14        "concept_choice": r["concept_choice"],
15        "core_integrations": r["core_integrations"],
16        "moral_integrations": r["moral_integrations"],
17        "underrepresented_integrations": r["underrepresented_integrations"],
18    }
19
20if __name__ == "__main__":
21
22    parser = argparse.ArgumentParser()
23    
24    parser.add_argument( "--data_dir",  type=str )
25    parser.add_argument( "--hub_path",  type=str )
26    parser.add_argument( "--n_test",    type=int )
27    parser.add_argument( "--rand_seed", type=int, default=32 )
28    
29    args = parser.parse_args()
30
31
32    data_dir = Path( args.data_dir )
33    n_test = args.n_test
34    rand_seed = args.rand_seed
35    hub_path = args.hub_path
36
37    # -------------------------------------------------------------------
38
39    data_file_paths = [ data_dir / f"data_{i}.json" for i in range( 1000 ) ]
40    data_file_paths = [ p for p in data_file_paths if p.exists() ] 
41
42    curated_rows = []
43    for data_path in data_file_paths :
44        data = load_json( data_path )
45        curated_rows.extend( [
46            curate_row( r ) for r in data[ "instances" ]
47        ] )
48
49    print(f"Total: {len(curated_rows)} rows")
50
51    # -------------------------------------------------------------------
52
53    ds = Dataset.from_list(curated_rows)
54    split = ds.train_test_split(test_size=n_test, seed=rand_seed )
55    dd = DatasetDict({"train": split["train"], "test": split["test"]})
56
57    dd.push_to_hub( hub_path )