radames's picture
Update modules/dataset.py
fc4c647
raw
history blame
366 Bytes
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