<jupyter_start><jupyter_text>Summary----  Imports<jupyter_code>import os
import time
from pathlib import Path

import tqdm

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq<jupyter_output><empty_output><jupyter_text>Parameters<jupyter_code>SUBSET = "training"
RANDOM_SEED = 0
DATAPKG_OUTPUT_DIR = Path("~/datapkg_output_dir").expanduser().resolve()
DATAPKG_OUTPUT_DIR
DEEP_PROTEIN_GEN_DIR = DATAPKG_OUTPUT_DIR / "deep-protein-gen"
DEEP_PROTEIN_GEN_DIR<jupyter_output><empty_output><jupyter_text>Workflow<jupyter_code>DATA_DIR = DEEP_PROTEIN_GEN_DIR / f"{SUBSET}_data"
DATA_DIR
OUTPUT_DATA_DIR = DEEP_PROTEIN_GEN_DIR / f"{SUBSET}_data_rs{RANDOM_SEED}"
OUTPUT_DATA_DIR
folder_list = os.listdir(DATA_DIR)
folder_list[:3]
columns = [
    "uniparc_id",
    "structure_id",
    "sequence",
    "residue_idx_1_corrected",
    "residue_idx_2_corrected",
    "distances",
]
columns
def write_table(table, output_file, max_tries=5):
    n_tries = 0
    while True:
        n_tries += 1
        try:
            pq.write_table(table, output_file)
            return
        except OSError:
            if n_tries > max_tries:
                raise
            else:
                time.sleep(5)
random_state = np.random.RandomState(RANDOM_SEED)

num_rows_per_group = 200

for folder in tqdm.tqdm_notebook(folder_list, total=len(folder_list)):
    file_list = list((DATA_DIR / folder).glob("*.parquet"))
    assert len(file_list) == 1

    database_id = folder.replace("database_id=G3DSA%3A", "")

    output_file = OUTPUT_DATA_DIR / f"{database_id}.parquet"
    if output_file.is_file():
        try:
            df = pq.read_table(output_file, use_threads=False).to_pandas(
                integer_object_nulls=True
            )
            assert len(df) > 0
            continue
        except Exception as e:
            print(type(e), e)

    pqfile = pq.ParquetFile(file_list[0])
    num_rows = pqfile.metadata.num_rows
    mask = random_state.rand(num_rows) < (
        num_rows_per_group / num_rows * max(1, num_rows / num_rows_per_group) ** 0.5
    )

    dfs = []
    start = 0
    for row_group in range(pqfile.num_row_groups):
        df = pqfile.read_row_group(
            row_group, use_threads=False, columns=columns
        ).to_pandas(integer_object_nulls=True)
        end = start + len(df)
        df = df.loc[mask[start:end], :]
        assert len(df.index) == len(df.index.drop_duplicates())
        dfs.append(df)
        start = end
    assert start == num_rows
    df = pd.concat(dfs, ignore_index=True, sort=False)
    df["database_id"] = database_id

    print(database_id, num_rows, mask.sum(), len(df))
    table = pa.Table.from_pandas(df, preserve_index=False)
    write_table(table, output_file)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1072838177251525, 'fraction_numerical': 0.007176175098672408, 'mean_word_length': 3.2695252679938744, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '185104', 'license': 'MIT', 'n_tokens_mistral': 1042, 'n_tokens_neox': 1011, 'n_words': 168, 'path': 'notebooks/02_prepare_dataset_manual.ipynb', 'repo_name': 'ostrokach/proteinsolver', 'size': '2787'}