- added support for jsonl.zst type when calling get_access_log()

modelscope · Dec 19, 2023 · 32f40dd · 32f40dd
1 parent 325a855
commit 32f40dd
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/data_juicer/utils/constant.py b/data_juicer/utils/constant.py
@@ -1,5 +1,6 @@
 import copy
 import inspect
+import io
 import os
 
 import zstandard as zstd
@@ -62,9 +63,9 @@ def get_access_log(cls, dj_cfg=None):
                         # Create a stream reader for the file and decode the
                         # first line
                         with dctx.stream_reader(compressed_file) as reader:
-                            first_line_bytes = reader.readline()
-                            # Assuming the file is encoded in UTF-8
-                            first_line = first_line_bytes.decode('utf-8')
+                            text_stream = io.TextIOWrapper(
+                                reader, encoding='utf-8')
+                            first_line = text_stream.readline()
                 elif 'jsonl' in dj_cfg.dataset_path:
                     tmp_f_name = dj_cfg.dataset_path. \
                         replace('.jsonl', '.tmp.jsonl')