add token mapping

delphi-suite · Feb 2, 2024 · 8339961 · 8339961
1 parent cb3976e
commit 8339961
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/scripts/.gitkeep b/scripts/.gitkeep
diff --git a/scripts/map_tokens.py b/scripts/map_tokens.py
@@ -0,0 +1,14 @@
+from datasets import DatasetDict, IterableDataset, IterableDatasetDict
+from datasets.arrow_dataset import Dataset
+
+
+def token_map(
+    tokenized_dataset: DatasetDict | Dataset | IterableDatasetDict | IterableDataset,
+) -> dict[int, list[tuple[int, int]]]:
+    mapping = {}
+
+    for prompt_idx, prompt in enumerate(tokenized_dataset):
+        for token_idx, token in enumerate(prompt["tokens"]):
+            mapping.setdefault(token, []).append((prompt_idx, token_idx))
+
+    return mapping