Sort labels (NVIDIA#221)

Signed-off-by: Ryan Wolf <[email protected]>
sarahyurick · Aug 29, 2024 · 2e0af57 · 2e0af57
1 parent cb2b705
commit 2e0af57
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/nemo_curator/classifiers/domain.py b/nemo_curator/classifiers/domain.py
@@ -95,6 +95,7 @@ def __init__(
 
         self.prob_column = prob_column
         self.labels = list(config.label2id.keys())
+        self.labels.sort(key=lambda x: config.label2id[x])
         self.out_dim = len(self.labels)
 
         model = DomainModel(

diff --git a/nemo_curator/classifiers/quality.py b/nemo_curator/classifiers/quality.py
@@ -92,6 +92,7 @@ def __init__(
 
         self.prob_column = prob_column
         self.labels = list(config.label2id.keys())
+        self.labels.sort(key=lambda x: config.label2id[x])
         self.out_dim = len(self.labels)
 
         model = QualityModel(