modules/dataset.py

from datasets import load_dataset

dataset = load_dataset("go_emotions", split="train")

emotions = dataset.info.features["labels"].feature.names


def query_emotion(start, end):
    rows = dataset[start:end]

    observations = [
        {"text": r[0], "emotion": emotions[r[1][0]]}
        for r in zip(rows["text"], rows["labels"])
    ]

    return observations