read and clean imdb data

read_clean_imdb_data.py

import pandas as pd
import re

# text and label column names
TEXT_COL = "text"
LABEL_COL = "label"

def clean_html(text: str):
    "remove html tags and whitespaces"
    cleanr = re.compile('<.*?>')
    text = re.sub(cleanr, '  ', text)
    return re.sub(' +', ' ', text)

def read_imdb(data_dir, max_lengths={"train": None, "test": None}):
    datasets = {}
    for t in ["train", "test"]:
        df = pd.read_csv(os.path.join(data_dir, f"imdb5k_{t}.csv"))
        if max_lengths.get(t) is not None:
            df = df.sample(n=max_lengths.get(t))
            df[TEXT_COL] = df[TEXT_COL].apply(lambda t: clean_html(t))
        datasets[t] = df
    return datasets    

# read data
datasets = read_imdb(IMDB_DIR)

# list of labels
labels = list(set(datasets["train"][LABEL_COL].tolist()))

# labels to integers mapping
label2int = {label: i for i, label in enumerate(labels)}