INNER CODE UNIT · Python

split_csv_by_max_len

CogStack/OpenGPT · opengpt/dataset_utils.py:12

def split_csv_by_max_len(datasets, max_len, tokenizer, base_path):
    r''' Given a tokenizer it will split the dataset (based on the `text` column) into max_len sequencse 
    '''
    for dataset in tqdm(datasets, desc='Datasets', total=len(datasets)):
        csv_path = dataset['path']
        name = dataset['name']

        nrows = None
        if dataset.get('nrows', -1) > 0:
            nrows = dataset['nrows']

        df = pd.read_csv(csv_path, nrows=nrows)
        cols = df.columns
        assert 'text' in cols, f'The CSV for dataset {name} has no "text" column.'

        new_data = [list(cols) + ['len', 'part']]
        for _, row in tqdm(df.iterrows(), desc=dataset['name'], total=len(df)):
            text = row['text']

View source record →