fix download wikipedia

specify number to embed, to skip

utils.py

@@ -1,4 +1,5 @@
 import os
+import re
 import time
 import shutil
 from pathlib import Path
@@ -107,33 +108,48 @@ def load_hf_dataset(ds_name: str, ds_config: str = None, ds_split: str = "train"
     if ds_config == "":
         ds_config = None
 
-    ds = load_dataset(ds_name, ds_config, split=ds_split, )
-    #streaming=True)
+    if ds_name == "wikipedia":
+        pattern = re.compile(r"[^a-zA-Z0-9]")
+        folder = Path("/data") / pattern.sub("", ds_name+ds_config)
+        files = list(map(str, folder.glob("chunk_")))
+        
+        return load_dataset("parquet", data_files=files, split="train")
+    
+    ds = load_dataset(ds_name, ds_config, split=ds_split)
 
     return ds
 
-def download_wikipedia(ds_name, ds_config):
+def download_wikipedia(ds_name, ds_config, num2skip, num2embed):
     ds = load_dataset(ds_name, ds_config, streaming=True, split="train")
 
     def gen():
-        for example in ds:
-            yield {"text": example["text"]}
+        if num2embed > 0:
+
+            for example in ds.skip(num2skip).take(num2embed):
+                yield {"text": example["text"]}
+        else:
+            for example in ds.skip(num2skip):
+                yield {"text": example["text"]}
 
     ds2 = Dataset.from_generator(gen)
 
-    chunk_size = 200_000
+    chunk_size = 20_000
 
     filenames = []
 
-    Path("wiki_chunks").mkdir(exist_ok=True)
+    pattern = re.compile(r"[^a-zA-Z0-9]")
+
+    folder = Path("/data") / pattern.sub("", ds_name+ds_config)
+
+    folder.mkdir(exist_ok=True, parents=True)
 
     for chunk_num, start_idx in enumerate(range(0, len(ds2), chunk_size)):
         end_idx = min(start_idx + chunk_size, len(ds2))
 
         temp = ds2.select(range(start_idx, end_idx))
 
-        temp.to_parquet(f"/data/wiki_chunks/chunk_{chunk_num}")
-        filenames.append(f"/data/wiki_chunks/chunk_{chunk_num}")
+        temp.to_parquet(str(folder / f"chunk_{chunk_num}"))
+        filenames.append(str(folder / f"chunk_{chunk_num}"))
 
     return load_dataset("parquet", data_files=filenames, split="train")