From d2d65c7efabe49573339a3790839311ab2b4f392 Mon Sep 17 00:00:00 2001
From: songsenand <songsenand@163.com>
Date: Sat, 14 Feb 2026 15:26:53 +0800
Subject: [PATCH] =?UTF-8?q?=E8=B0=83=E6=95=B4=E5=AF=BC=E5=85=A5=E9=A1=BA?=
 =?UTF-8?q?=E5=BA=8F=E5=B9=B6=E4=BF=AE=E5=A4=8Dpickle=E4=BF=9D=E5=AD=98?=
 =?UTF-8?q?=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/tmp_utils/gen_eval_dataset.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/tmp_utils/gen_eval_dataset.py b/src/tmp_utils/gen_eval_dataset.py
index 24a1852..d13b4c3 100644
--- a/src/tmp_utils/gen_eval_dataset.py
+++ b/src/tmp_utils/gen_eval_dataset.py
@@ -1,13 +1,13 @@
-from tqdm import tqdm
-from loguru import logger
-import torch
-from torch.utils.data import DataLoader
 import pickle
 from pathlib import Path
 
-from suinput.dataset import PinyinInputDataset, worker_init_fn, custom_collate_with_txt
-from suinput.query import QueryEngine
+import torch
+from loguru import logger
+from torch.utils.data import DataLoader
+from tqdm import tqdm
 
+from suinput.dataset import PinyinInputDataset, custom_collate_with_txt, worker_init_fn
+from suinput.query import QueryEngine
 
 # 使用示例
 if __name__ == "__main__":
@@ -42,7 +42,13 @@ if __name__ == "__main__":
         for i, sample in tqdm(enumerate(dataloader), total=5):
             if i >= total:
                 break
-            print(sample)
-            # pickle.dump(sample, open(f"{str(Path(__file__).parent.parent / 'trainer' / 'eval_dataset')}/sample_{i}.pkl", "wb"))
+            # print(sample)
+            pickle.dump(
+                sample,
+                open(
+                    f"{str(Path(__file__).parent.parent / 'trainer' / 'eval_dataset')}/sample_{i}.pkl",
+                    "wb",
+                ),
+            )
     except StopIteration:
         print("数据集为空")