File size: 3,538 Bytes
a7ef49b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
#./experiments/experiment1/preprocessing.py
import logging
import os
import sqlite3
from transformers import GPT2TokenizerFast
from datasets import load_dataset
class DatabaseInterface(object):
def __init__(self, db_file):
self.db_file = db_file
def create_table(self, table_name=None):
conn = sqlite3.connect(self.db_file)
c = conn.cursor()
c.execute(
'''
CREATE TABLE IF NOT EXISTS plain_text (
text TEXT,
split TEXT
)
'''
)
conn.commit()
conn.close()
def write_plain_text(self, example, split):
conn = sqlite3.connect(self.db_file)
c = conn.cursor()
c.execute("INSERT INTO plain_text (text, split) VALUES (?, ?)",
(example, split))
conn.commit()
conn.close()
def process_and_write(example, writer, split):
writer.write_plain_text(example, split)
def prepare_data(start_index, end_index, **kwargs):
data_writer = kwargs['data_writer']
train_dataset = kwargs['train_dataset']
val_dataset = kwargs['val_dataset']
for split, dataset in {'val': val_dataset, 'train': train_dataset}.items():
subset = dataset[start_index:end_index] # Select the subset based on start and end indices
if isinstance(subset, dict):
subset = subset["text"] # Extract the "text" part from the subset dictionary
for example in subset:
process_and_write(example, data_writer, split)
if __name__ == '__main__':
logging.basicConfig(
format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO
)
# Configs
batch_size = 32
num_processes = 4 # number of jobs to run simultaneously
logging.info("Creating Database Interface")
db_file_path = os.path.join('data', 'experiment1.db')
_delete_db = True
# Check to see if the database file already exists
if os.path.exists(db_file_path):
if _delete_db:
logging.info(f"Database file {db_file_path} already exists. Deleting it.")
os.remove(db_file_path)
data_writer = DatabaseInterface(db_file_path)
data_writer.create_table()
logging.info("Database table `plain_text` created")
else:
logging.info(f"Database file {db_file_path} already exists. Connecting to it.")
data_writer = DatabaseInterface(db_file_path)
else:
data_writer = DatabaseInterface(db_file_path)
data_writer.create_table()
logging.info("Database table `plain_text` created")
#cache_dir=os.path.join(
# 'C:/Users/User/.cache/huggingface/datasets/openwebtext/plain_text',
# '1.0.0',
# '6f68e85c16ccc770c0dd489f4008852ea9633604995addd0cd76e293aed9e521'
#)
dataset = load_dataset(
"openwebtext",
cache_dir=cache_dir,
num_proc=num_processes,
save_infos = True,
writer_batch_size=batch_size
)
split_dataset = dataset["train"].train_test_split(test_size=0.1, seed=42, shuffle=False)
train_dataset = split_dataset["train"]
val_dataset = split_dataset["test"]
prepare_data(
start_index=0,
end_index=1000,
**{
'data_writer': data_writer,
'train_dataset': train_dataset,
'val_dataset': val_dataset,
}
)
|