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]
    texts, labels = [rows[k] for k in rows.keys()]

    observations = []

    for i, text in enumerate(texts):
        observations.append({
            "text": text,
            "emotion": emotions[labels[i]],
        })

    return observations