data_prep.py

from datasets import load_dataset
from utils import save_data
import re

def clean_lines(lines):
    """
    Funtion for cleaining the data, removes all dubble spaces and all tabs
    """
    out = []
    for s in lines.splitlines():
        s = re.sub(r"\s+", " ", s).strip()
        #s = s.replace("\t\t", " ").strip()
        if s:
            out.append(s)
    return "\n".join(out) 

if __name__ == "__main__":

    out_src_train = "data/wikitext2_train.txt"
    out_src_valid = "data/wikitext2_valid.txt"
    out_src_test = "data/wikitext2_test.txt"

    # Download Wikitext2
    ds = load_dataset("wikitext", "wikitext-2-raw-v1")  # 'train'/'validation'/'test'
    
    # Split into train, val and text
    text_train = "\n\n".join(ds["train"]["text"])
    text_valid = "\n\n".join(ds["validation"]["text"])
    text_test = "\n\n".join(ds["test"]["text"])

    # Clean the data
    text_train_cl = clean_lines(text_train)
    text_valid_cl = clean_lines(text_valid)
    text_test_cl = clean_lines(text_test)

    save_data(out_src=out_src_train, text=text_train_cl)
    save_data(out_src=out_src_test, text=text_test_cl)
    save_data(out_src=out_src_valid, text=text_valid_cl)