Make mypy happy

flairNLP · Dec 17, 2023 · f180771 · f180771
1 parent d0441d7
commit f180771
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/flair/datasets/document_classification.py b/flair/datasets/document_classification.py
@@ -965,11 +965,10 @@ def __init__(
                         label, title, description = row
                         # Original labels are [1, 2, 3, 4] -> ['World', 'Sports', 'Business', 'Sci/Tech']
                         # Re-map to [0, 1, 2, 3].
-                        label = int(label) - 1
                         text = " ".join((title, description))
 
                         new_label = "__label__"
-                        new_label += label_dict[label]
+                        new_label += label_dict[int(label) - 1]
 
                         write_fp.write(f"{new_label} {text}\n")