calculate_perplexity.py

import datasets
import kenlm

kenlm_model = kenlm.Model("PATH_TO_YOUR_KENLM")

def get_perplexity(example):
    """Get perplexity of text field."""
    example["perplexity_kenlm_full"] = kenlm_model.perplexity(example["text"].lower())
    return example

dataset = datasets.load_dataset("csv", data_files={"train": "train.csv", "validation": "valid.csv"})

# TRAIN SPLIT PERPLEXITY

dataset["train"] = dataset["train"].map(get_perplexity, num_proc=64, writer_batch_size=100000)

# VALIDATION SPLIT PERPLEXITY

dataset["validation"] = dataset["validation"].map(get_perplexity, num_proc=64, writer_batch_size=100000)

# SAVE DATASET

dataset["train"].to_csv("train.csv", num_proc=64, index=False)
dataset["validation"].to_csv("valid.csv", num_proc=64, index=False)