Fixed bug in #131

nicolay-r · Jul 23, 2022 · c268ec6 · c268ec6
1 parent 165f822
commit c268ec6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/arekit/contrib/source/rusentrel/io_utils.py b/arekit/contrib/source/rusentrel/io_utils.py
@@ -55,8 +55,8 @@ def __get_root_by_index(doc_id, version, keep_etalon=False):
         assert(RuSentRelIOUtils.__is_supported(version))
         assert(isinstance(version, RuSentRelVersions))
         assert(isinstance(doc_id, int))
-        other_dir = RuSentRelIOUtils.ETALON_FOLDER if keep_etalon else RuSentRelIOUtils.TRAIN_FOLDER
-        test_indices = set(RuSentRelIOUtils.__iter_indicies_from_dataset(version, RuSentRelIOUtils.TRAIN_FOLDER))
+        other_dir = RuSentRelIOUtils.ETALON_FOLDER if keep_etalon else RuSentRelIOUtils.TEST_FOLDER
+        test_indices = set(RuSentRelIOUtils.__iter_indicies_from_dataset(version, RuSentRelIOUtils.TEST_FOLDER))
         return other_dir if doc_id in test_indices else RuSentRelIOUtils.TRAIN_FOLDER
 
     @staticmethod