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")