Fix score call indentation

Signed-off-by: Ryan Wolf <[email protected]>
NVIDIA · Mar 27, 2024 · 34f2d7f · 34f2d7f
1 parent 2e96a9c
commit 34f2d7f
Showing 1 changed file with 17 additions and 17 deletions.
diff --git a/nemo_curator/modules/filter.py b/nemo_curator/modules/filter.py
@@ -31,23 +31,23 @@ def __init__(self, score_fn, score_field, text_field="text", score_type=None):
         self.text_field = text_field
         self.score_type = score_type
 
-        def __call__(self, dataset):
-            # Set the metadata for the function calls if provided
-            if self.score_type:
-                meta = (None, self.score_type)
-            else:
-                meta = no_default
-
-            if is_batched(self.score_fn):
-                dataset.df[self.score_field] = dataset.df[
-                    self.text_field
-                ].map_partitions(self.score_fn, meta=meta)
-            else:
-                dataset.df[self.score_field] = dataset.df[self.text_field].apply(
-                    self.score_fn, meta=meta
-                )
-
-            return dataset
+    def __call__(self, dataset):
+        # Set the metadata for the function calls if provided
+        if self.score_type:
+            meta = (None, self.score_type)
+        else:
+            meta = no_default
+
+        if is_batched(self.score_fn):
+            dataset.df[self.score_field] = dataset.df[self.text_field].map_partitions(
+                self.score_fn, meta=meta
+            )
+        else:
+            dataset.df[self.score_field] = dataset.df[self.text_field].apply(
+                self.score_fn, meta=meta
+            )
+
+        return dataset
 
 
 class Filter: