test

Tencent · Dec 25, 2023 · d1af253 · d1af253
1 parent 51350e6
commit d1af253
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/tencentpretrain/utils/dataset.py b/tencentpretrain/utils/dataset.py
@@ -247,10 +247,12 @@ def worker(self, proc_id, start, end):
                     if self.json_format_corpus:
                         data = json.loads(line)
                         line = data.get("text", "") + data.get("content", "")
+                        if len(line) < 5:
+                            continue
+                        if pos >= end:
+                            break
 
                     pos += 1
-                    if len(line) < 5:
-                        continue
 
                     document = [self.vocab.get(CLS_TOKEN)] + self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(line)) + [self.vocab.get(SEP_TOKEN)]
 
@@ -281,8 +283,7 @@ def worker(self, proc_id, start, end):
                             for instance in instances:
                                 pickle.dump(instance, dataset_writer)
 
-                    if pos >= end:
-                        break
+
 
         dataset_writer.close()