Update dataloader.py

Shark-NLP · May 10, 2023 · 78d88de · 78d88de
1 parent 9876af6
commit 78d88de
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/model/dataloader.py b/model/dataloader.py
@@ -1,8 +1,9 @@
 from fastNLP.io import JsonLoader
 from transformers import AutoTokenizer
+from preprocess.preprocess import TOKENIED_FILE_DIR
 
-def get_data_path(PTM, dataset):
-    paths = {'train': f'tokenized_files/{dataset}/train.{PTM}.jsonl', 'val': f'tokenized_files/{dataset}/val.{PTM}.jsonl'}
+def get_data_path(PTM, dataset):   
+    paths = {'train': f'{TOKENIED_FILE_DIR}/{dataset}/train.{PTM}.jsonl', 'val': f'{TOKENIED_FILE_DIR}/{dataset}/val.{PTM}.jsonl'}
     return paths