datasets.map(), datasets.filter()の仕様について

TL; DR

背景

アルバイトの研修タスクに取り組んでいる途中、Hugging Faceのdatasets.map()関数の中身の書き方がわからなくなり、NLP Course 5章に取り組みながら解決したので忘備録として残します。

その時使っていたデータセット: drug review dataset
)

!wget "https://archive.ics.uci.edu/ml/machine-learning-databases/00462/drugsCom_raw.zip"
!unzip drugsCom_raw.zip

from datasets import load_dataset

data_files = {"train": "drugsComTrain_raw.tsv", "test": "drugsComTest_raw.tsv"}
drug_dataset = load_dataset("csv", data_files=data_files, delimiter="\t")

わかったこと

filter()メソッドでデータセットのフィルタリングをしたいとき：
条件をラムダ関数などで記述して渡す

drug_dataset = drug_dataset.filter(lambda x: x["condition"] is not None)

既存の列の値に変更を加えたいとき：
処理したい列のみ辞書形式でreturnすれば処理結果がデータセットに反映される

def lowercase_condition(example):
    return {"condition": example["condition"].lower()}

drug_dataset = drug_dataset.map(lowercase_condition)

既存の列の値から新しい列を生成したいとき:
新しい列を作りたい際も辞書形式で新しい列名のみ書いてリターンすればデータセットに反映される

def compute_review_length(example):
    return {"review_length": len(example["review"].split())}

drug_dataset = drug_dataset.map(compute_review_length)