data_utils.py

import re
from datasets import load_dataset, Dataset, DatasetDict
from itertools import chain
from tqdm import tqdm
from collections import Counter
from accelerate import Accelerator

LANGUAGES_TO_DECODE_FROM_BYTES = ["he", "fr", "uk"]
STREAMING_DATASETS = ["fineweb-edu"]


def load_pg19_val_and_test():
    # Load the dataset in streaming mode
    streaming_dataset = load_dataset("deepmind/pg19", split=None, streaming=True)

    # Extract test and validation splits
    test_split = list(streaming_dataset["test"])
    validation_split = list(streaming_dataset["validation"])

    # Convert them into regular datasets
    test_dataset = Dataset.from_list(test_split)
    validation_dataset = Dataset.from_list(validation_split)

    # validation_dataset = load_dataset("deepmind/pg19", split="validation")
    # test_dataset = load_dataset("deepmind/pg19", split="test")

    return DatasetDict({"validation": validation_dataset, "test": test_dataset})


def load_pubmed(n_samples=10000):
    # Load the dataset in streaming mode
    streaming_dataset = load_dataset("MedRAG/pubmed", streaming=True)

    # Extract test and validation splits
    data = list(streaming_dataset["train"].take(n_samples*4))
    train = data[:2*n_samples]
    validation = data[2*n_samples:3*n_samples]
    test = data[3*n_samples:]
    # Convert them into regular datasets
    train = Dataset.from_list(train)
    validation = Dataset.from_list(validation)
    test = Dataset.from_list(test)
    dataset = DatasetDict({"train": train, 'validation': validation, 'test': test})
    dataset = dataset.rename_column('content', 'text')
    return dataset


def load_lm_dataset(dataset_name, language="en", split=None):
    """
    Loads a popular pretraining or perplexity evaluation dataset by name and language.

    Args:
        dataset_name (str): The name of the dataset to load. Options include:
            - 'wikitext' (wikitext-2, smaller WikiText dataset)
            - 'wikitext-103' (larger WikiText dataset)
            - 'pg19' (Project Gutenberg dataset for long-context modeling)
            - 'c4' (Common Crawl-based English corpus)
            - 'wiki40b' (Wikipedia dataset in multiple languages)
            - 'mc4' (Multilingual C4 dataset in various languages)
        language (str): Language code for datasets that support multilingual options (e.g., 'en' for English).
                        Defaults to 'en'.

    Returns:
        Dataset: Loaded Hugging Face dataset.
    """
    if dataset_name.lower() == 'wikitext':
        return load_dataset("Salesforce/wikitext", "wikitext-2-raw-v1", split=split)
    elif dataset_name.lower() == 'fineweb-edu':
        return load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT")
    elif dataset_name.lower() == 'wikitext-103':
        return load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split=split)
    elif dataset_name.lower() == 'cord19':
        return load_dataset("allenai/cord19", "fulltext", trust_remote_code=True)
    elif dataset_name.lower() == 'pubmed':
        return load_pubmed()
    elif dataset_name.lower() == 'wikilingua':
        dataset = load_dataset("GEM/wiki_lingua", trust_remote_code=True)
        dataset = dataset.filter(lambda ex: (ex['source_language'] == "en") & (ex['target_language'] == "en"))
        dataset = dataset.rename_column("source", "text")
        dataset = dataset.rename_column("target", "summary")
        return dataset
    elif dataset_name.lower() == 'xsum':
        dataset = load_dataset("EdinburghNLP/xsum")
        dataset = dataset.rename_column("document", "text")
        return dataset
    elif dataset_name.lower() == 'cnn':
        dataset = load_dataset("abisee/cnn_dailymail", "3.0.0")
        dataset = dataset.rename_column("article", "text")
        dataset = dataset.rename_column("highlights", "summary")
        dataset = dataset.map(lambda example: {"text": example["text"].replace("(CNN)", "")})
        return dataset
    elif dataset_name.lower() == 'pg19':
        return load_pg19_val_and_test()
    elif dataset_name.lower() == 'wiki40b':
        dataset = load_dataset("google/wiki40b", language, split=split)
        if language in LANGUAGES_TO_DECODE_FROM_BYTES:
            dataset = dataset.map(lambda x: {
                "text": bytes(x["text"][2:-1], "utf-8").decode("unicode_escape").encode("latin1").decode("utf-8").replace("_NEWLINE_", "\n")
            })
        return dataset
    else:
        raise ValueError(
            "Dataset not recognized. Available options: 'wikitext-2', 'wikitext-103', 'pg19', 'c4', 'wiki40b', 'mc4'.")


def extract_new_words_from_dataset(
        dataset: Dataset, tokenizer, text_column: str = "text", max_samples: int = None, filter_func=(lambda word, token_count: True)):
    """
    Loads a Hugging Face dataset and extracts all unique words from the specified text column.

    Args:
        dataset (Dataset): Name of the dataset to load.
        split (str): Dataset split to use, typically 'train' for training data. Defaults to 'train'.
        text_column (str): The column in the dataset containing text. Defaults to 'text'.
        max_samples (int): Number of samples from the dataset to go over.

    Returns:
        set: A set of unique words in the dataset.
    """
    if max_samples:
        dataset = dataset.select(range(max_samples))

    # Regular expression to split text into words (adjust as needed for specific languages)
    # word_pattern = re.compile(r"\b\w+\b")
    word_pattern = re.compile(r"\b\w+(?:[-']\w+)*\b")

    # Iterate over each entry in the dataset and extract unique words
    all_words = list()
    new_words = list()
    for record in tqdm(dataset, total=len(dataset), miniters=10, desc="Extracting all words from dataset...", unit="examples"):
        text = record.get(text_column, "")
        words = word_pattern.findall(text)
        all_words += words

    # all_words = list(dict.fromkeys(all_words))
    word_frequencies = Counter(all_words)
    all_words = list(word_frequencies.keys())
    token_counts = [len(x) for x in tokenizer(all_words, add_special_tokens=False)["input_ids"]]
    w_whitespace_token_counts = [len(x) for x in tokenizer([f" {w}" for w in all_words], add_special_tokens=False)["input_ids"]]

    new_words = [word for word, count, w_whitespace_count in zip(all_words, token_counts, w_whitespace_token_counts) if ((count > 1) and (w_whitespace_count > 1) and filter_func(word, count))]
    new_words_freq = {word: word_frequencies[word] for word in new_words}
    # for word, token_count in tqdm(all_words, total=len(all_words), miniters=10, desc="Finding new words...", unit="words"):
    #     if (not tokenizer.vocab.get(word, False)) and :
    #         new_words.append(word)

    # remove duplicates and return
    return new_words, new_words_freq


def get_group_texts_func(block_size=1024):
    def group_texts(examples):
        # Concatenate all texts.
        concatenated_examples = {k: list(chain(*examples[k])) for k in examples.keys()}

        total_length = len(concatenated_examples[list(examples.keys())[0]])
        # We drop the small remainder, and if the total_length < block_size  we exclude this batch and return an empty dict.
        # We could add padding if the model supported it instead of this drop, you can customize this part to your needs.
        total_length = (total_length // block_size) * block_size
        # Split by chunks of max_len.
        result = {
            k: [t[i: i + block_size] for i in range(0, total_length, block_size)]
            for k, t in concatenated_examples.items()
        }
        result["labels"] = result["input_ids"].copy()
        return result
    return group_texts


def get_tokenize_func(tokenizer, text_col_name):
    def _tokenize(examples):
        output = tokenizer(
            examples[text_col_name],
            return_token_type_ids=False,
            add_special_tokens=False,
        )
        return output
    return _tokenize


def tokenize_and_prepare_dataset(
        dataset, tokenizer, accelerator=None,
        text_col_name: str = "text",
        max_length: int = 256,
        eval_max_samples: int = None,
):

    if tokenizer.bos_token is not None and max_length:
        # leave room for <BOS> token to be added:
        max_tokenized_len = max_length - 1
    else:
        max_tokenized_len = max_length

    tokenize_function = get_tokenize_func(tokenizer, text_col_name)

    column_names = dataset.column_names

    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=column_names,
        load_from_cache_file=False,
        desc="Running tokenizer on dataset",
    )
    group_texts = get_group_texts_func(block_size=max_tokenized_len)
    lm_dataset = tokenized_dataset.map(
        group_texts,
        batched=True,
    )

    if eval_max_samples:
        lm_dataset = lm_dataset.select(range(eval_max_samples))

    return lm_dataset