From b508402724667ca89e52c256891886fceda84bb9 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 30 Jun 2025 22:00:15 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=A4=84=E7=90=86=E5=BD=A2?=
 =?UTF-8?q?=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 to_ids.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/to_ids.py b/to_ids.py
index 0a7c447..2cf51a8 100644
--- a/to_ids.py
+++ b/to_ids.py
@@ -12,23 +12,26 @@ def fetch_files(directory):
 def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path):
     arrows = []
     with open(file_path, "r") as f:
-        for line in f:
+        lines = f.readlines()
+        file_name = os.path.basename(file_path)
+        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
             line = json.loads(line)
             ids = tokenizer.encode(line["text"])
-            arrows.append(ids)
+            arrow = torch.tensor(ids, dtype=torch.int32)
+            arrows.append(arrow)
     
     with open(out_file_path, "w") as f:
-        tensor = torch.tensor(arrows, dtype=torch.int32)
+        tensor = torch.cat(arrows)
         pkl.dump(tensor, f)
         
 def process_files(tokenizer: BpeTokenizer, files: List[str], base_out_dir):
-    for file_path in tqdm(files, desc="Processing files", total=len(files)):
+    for file_path in files:
         out_file_name = os.path.basename(file_path).replace(".jsonl", ".pkl")
         out_file_path = os.path.join(base_out_dir, out_file_name)
         
         if not os.path.exists(out_file_path):
             os.makedirs(os.path.dirname(out_file_path), exist_ok=True)
-            
+         
         convert_to_ids(tokenizer, file_path, out_file_path)