🐛 Fixed TypeError when printing DatasetWithEntities (#4)

Signed-off-by: Marcos Martinez <Marcos.Martinez.Galindo@ibm.com>
IBM · Sep 28, 2022 · cb7af04 · cb7af04
1 parent bf4fbe8
commit cb7af04
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/zshot/evaluation/dataset/dataset.py b/zshot/evaluation/dataset/dataset.py
@@ -30,4 +30,4 @@ def from_dict(
 
     def __repr__(self):
         return f"Dataset({{\n    features: {list(self.features.keys())},\n    num_rows: {self.num_rows}," \
-               f"\n    entities: {[ent['label'] for ent in self.entities if self.entities is not None]}\n}})"
+               f"\n    entities: {[ent.name for ent in self.entities if self.entities is not None]}\n}})"