Disable string conversion globally (#56)

Signed-off-by: Ryan Wolf <[email protected]>
NVIDIA · May 7, 2024 · 0e9b7f1 · 0e9b7f1
1 parent 8dd4255
commit 0e9b7f1
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/config/fasttext_langid.yaml b/config/fasttext_langid.yaml
@@ -1,5 +1,6 @@
 input_field: text
 filters:
   - name: nemo_curator.filters.classifier_filter.FastTextLangId
+    log_score: True
     params:
       model_path: <Path to the FasText language id model (e.g., lid.176.bin)>
diff --git a/nemo_curator/__init__.py b/nemo_curator/__init__.py
@@ -12,4 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dask
+
 from .modules import *
+
+# Dask will automatically convert the list score type
+# to a string without this option.
+# See https://github.com/NVIDIA/NeMo-Curator/issues/33
+# This also happens when reading and writing to files
+dask.config.set({"dataframe.convert-string": False})
diff --git a/nemo_curator/filters/classifier_filter.py b/nemo_curator/filters/classifier_filter.py
@@ -76,11 +76,6 @@ def __init__(self, model_path=None, min_langid_score=0.3):
         self._cutoff = min_langid_score
         self._name = "lang_id"
 
-        # Dask will automatically convert the list score type
-        # to a string without this option.
-        # See https://github.com/NVIDIA/NeMo-Curator/issues/33
-        dask.config.set({"dataframe.convert-string": False})
-
     @batched
     def score_document(self, df: pd.Series):
         model_attr = f"{self._name}_{self._model_path}"