From a8f05cccd6528f13171cd91fccb79a6a76b93364 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Mon, 30 Jun 2025 22:36:09 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E8=BE=93=E5=87=BA=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6=E7=9A=84bug?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .gitignore | 1 +
 to_ids.py  | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index b2426b8..4204de1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@ __pycache__*
 
 # dataset
 dataset/*
+pkl_output/*
 
 # tokenzier
 tokenizer.json
diff --git a/to_ids.py b/to_ids.py
index 2cf51a8..62aa857 100644
--- a/to_ids.py
+++ b/to_ids.py
@@ -20,7 +20,7 @@ def convert_to_ids(tokenizer: BpeTokenizer, file_path, out_file_path):
             arrow = torch.tensor(ids, dtype=torch.int32)
             arrows.append(arrow)
     
-    with open(out_file_path, "w") as f:
+    with open(out_file_path, "wb") as f:
         tensor = torch.cat(arrows)
         pkl.dump(tensor, f)
         
@@ -41,7 +41,7 @@ if __name__ == "__main__":
         os.path.join("dataset", "chinese-c4"),
         os.path.join("dataset", "english-fineweb")
     ]
-    base_out_dir = "cache"
+    base_out_dir = "pkl_output"
     
     files = []
     for dir_path in base_dir: