Update modules/dataset.py

modules/dataset.py

@@ -2,18 +2,15 @@ from datasets import load_dataset
 
 dataset = load_dataset("go_emotions", split="train")
 
-emotions = dataset.info.features['labels'].feature.names
+emotions = dataset.info.features["labels"].feature.names
+
 
 def query_emotion(start, end):
     rows = dataset[start:end]
-    texts, labels, id = [rows[k] for k in rows.keys()]
-
-    observations = []
 
-    for i, text in enumerate(texts):
-        observations.append({
-            "text": text,
-            "emotion": emotions[labels[i]],
-        })
+    observations = [
+        {"text": r[0], "emotion": emotions[r[1][0]]}
+        for r in zip(rows["text"], rows["labels"])
+    ]
 
     return observations