import pandas as pd
from pathlib import Path
from os.path import join
from tqdm import tqdm

if __name__ == '__main__':

    config = {
        'path_prefix': '/mnt/home/dongkeun/L2U/',
        'file': 'outputs/125m_train_init_512.csv',
        'doc_id': ['3409', '4983', '2302'],
    }
    df = pd.read_csv(join(config['path_prefix'], config['file']))

    for id_ in tqdm(config['doc_id']):
        text = df.loc[df['doc_id'] == int(id_)].iloc[0]['text']
        original = pd.DataFrame({'text': [text] * 1})

        Path(join(config['path_prefix'], f'data/injects')).mkdir(parents=True, exist_ok=True)
        original.to_csv(join(config['path_prefix'], f'data/injects/{id_}.csv'), index=False)