store locally

utils.py

@@ -19,7 +19,6 @@ from optimum.onnxruntime import (
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-
 opt_configs = {
     "O2": AutoOptimizationConfig.O2(),
     "O3": AutoOptimizationConfig.O3(),
@@ -108,7 +107,8 @@ 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)
+    ds = load_dataset(ds_name, ds_config, split=ds_split, )
+    #streaming=True)
 
     return ds
 
@@ -212,22 +212,34 @@ def tokenize(
     )
 
 
-def collate_fn(examples, tokenizer=None, padding=None, device=None):
+def collate_fn(examples, tokenizer=None, padding=None, column_name="text"):
     try:
         keys = examples[0].keys()
     except KeyError:
         print(examples)
     else:
         batch = {k: [] for k in examples[0].keys()}
+
+    tokenized = tokenizer(
+        [x[column_name] for x in examples],
+        truncation=True, 
+        padding=padding, 
+        max_length=512,
+        return_tensors="pt"
+        )
+
+    tokenized[column_name] = [x[column_name] for x in examples]
+
+    return tokenized
         
 
-    for example in examples:
-        for k, v in example.items():
-            batch[k].append(v)
+    # for example in examples:
+    #     for k, v in example.items():
+    #         batch[k].append(v)
 
-    return {
-        k: torch.tensor(v, dtype=torch.long) if k in {"attention_mask", "input_ids"} else v for k, v in batch.items()
-    }
+    # return {
+    #     k: torch.tensor(v, dtype=torch.long) if k in {"attention_mask", "input_ids"} else v for k, v in batch.items()
+    # }
 
 @torch.inference_mode()
 def batch_embed(
@@ -293,16 +305,16 @@ def batch_embed(
 
     repo = init_git_repo(new_dataset_id)
 
-    ds = ds.map(
-        tokenize,
-        batched=True,
-        batch_size=map_batch_size,
-        fn_kwargs={
-            "tokenizer": tokenizer,
-            "column_name": column_name,
-            "padding": "max_length" if opt_level == "O4" else True,
-        },
-    )
+    # ds = ds.map(
+    #     tokenize,
+    #     batched=True,
+    #     batch_size=map_batch_size,
+    #     fn_kwargs={
+    #         "tokenizer": tokenizer,
+    #         "column_name": column_name,
+    #         "padding": "max_length" if opt_level == "O4" else True,
+    #     },
+    # )
 
     embeds = []
     texts = []
@@ -327,10 +339,15 @@ def batch_embed(
         ds,
         batch_size=inference_bs,
         shuffle=False,
-        num_workers=1,
+        num_workers=2,
         pin_memory=True,
         drop_last=False,
-        collate_fn=partial(collate_fn, device=device)
+        collate_fn=partial(
+            collate_fn, 
+            column_name=column_name, 
+            tokenizer=tokenizer,
+            padding="max_length" if opt_level == "O4" else True
+            )
     ):
         ids = batch["input_ids"].to(device)
         mask = batch["attention_mask"].to(device)
@@ -354,7 +371,7 @@ def batch_embed(
 
         # Periodically upload to the hub
         if len(embeds) > upload_batch_size:
-            push_to_repo(repo, last_count, current_count, embeds, texts, api)
+            push_to_repo(new_dataset_id, last_count, current_count, embeds, texts, api)
             embeds = []
             texts = []
             last_count = current_count
@@ -372,7 +389,7 @@ def batch_embed(
 
     # If there are any remaining embeddings, upload them
     if len(embeds) > 0:
-        push_to_repo(repo, last_count, current_count, embeds, texts, api)
+        push_to_repo(new_dataset_id, last_count, current_count, embeds, texts, api)
 
     return current_count - num2skip, time_taken
 
@@ -472,27 +489,15 @@ def push_to_repo(
     files = sorted(list(data_dir.glob("*.parquet")))
 
 
-    if len(files) == 1:
-        api.upload_folder(
-            folder_path=str(data_dir),
-            repo_id=repo_id,
-            repo_type="dataset",
-            run_as_future=True,
-            token=os.environ["HF_TOKEN"],
-            commit_message=f"Embedded examples {last_count} thru {current_count} with folder",
-        )
-
-    else:
-
-        api.upload_file(
-            path_or_fileobj=filepath,
-            path_in_repo=f"data/{filename}",
-            repo_id=repo_id,
-            repo_type="dataset",
-            run_as_future=True,
-            token=os.environ["HF_TOKEN"],
-            commit_message=f"Embedded examples {last_count} thru {current_count}",
-        )
+    api.upload_file(
+        path_or_fileobj=filepath,
+        path_in_repo=f"data/{filename}",
+        repo_id=repo_id,
+        repo_type="dataset",
+        run_as_future=True,
+        token=os.environ["HF_TOKEN"],
+        commit_message=f"Embedded examples {last_count} thru {current_count}",
+    )
 
 
     # Delete old files