radames commited on
Commit
fc4c647
1 Parent(s): 9e0e6dd

Update modules/dataset.py

Browse files
Files changed (1) hide show
  1. modules/dataset.py +6 -9
modules/dataset.py CHANGED
@@ -2,18 +2,15 @@ from datasets import load_dataset
2
 
3
  dataset = load_dataset("go_emotions", split="train")
4
 
5
- emotions = dataset.info.features['labels'].feature.names
 
6
 
7
  def query_emotion(start, end):
8
  rows = dataset[start:end]
9
- texts, labels, id = [rows[k] for k in rows.keys()]
10
-
11
- observations = []
12
 
13
- for i, text in enumerate(texts):
14
- observations.append({
15
- "text": text,
16
- "emotion": emotions[labels[i]],
17
- })
18
 
19
  return observations
 
2
 
3
  dataset = load_dataset("go_emotions", split="train")
4
 
5
+ emotions = dataset.info.features["labels"].feature.names
6
+
7
 
8
  def query_emotion(start, end):
9
  rows = dataset[start:end]
 
 
 
10
 
11
+ observations = [
12
+ {"text": r[0], "emotion": emotions[r[1][0]]}
13
+ for r in zip(rows["text"], rows["labels"])
14
+ ]
 
15
 
16
  return observations