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()