process_the_pile_partition.py

import zstandard as zstd
import json
from tqdm.notebook import tqdm
from datasets import load_dataset

pile_path = "the_pile/train/00.jsonl.zst"

arxiv = []
with zstd.open(pile_path, 'r') as f:
    for i, line in enumerate(tqdm(f)):
        doc = json.loads(line)
        source = doc['meta']['pile_set_name']
        if source == "ArXiv":
            arxiv.append(doc)
            
with open("/tmp/arxiv.jsonl", "w") as f:
    for entry in arxiv:
        json.dump(entry, f)
        f.write("\n")
        
        
pile_arxiv = load_dataset("json", data_files={'train': "/tmp/arxiv.jsonl", 'validation': "/tmp/arxiv_val.jsonl", 'test': "/tmp/arxiv_test.jsonl"})

pile_arxiv.push_to_hub("haritzpuerto/the_pile_arxiv")