bugfix outputfilename with gz

huggingface · Dec 12, 2023 · 476de37 · 476de37
1 parent 35fc53e
commit 476de37
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/datatrove/pipeline/writers/jsonl.py b/src/datatrove/pipeline/writers/jsonl.py
@@ -1,5 +1,6 @@
 import dataclasses
 import json
+from string import Template
 
 from datatrove.data import Document
 from datatrove.io import BaseOutputDataFile, BaseOutputDataFolder
@@ -14,7 +15,7 @@ def __init__(self, output_folder: BaseOutputDataFolder, output_filename: str = N
  super().__init__(output_folder, output_filename=output_filename)
  self.gzip = gzip
  if self.gzip:
- self.output_filename = self.output_filename + ".gz"
+ self.output_filename = Template(self.output_filename.template + ".gz")
 
  def open(self, output_filename):
  return self.output_folder.open(output_filename, mode="wt", gzip=self.gzip)