(
self, root, split="train", break_mode="none", max_length=256, min_length=1
)
| 61 | """ |
| 62 | |
| 63 | def __init__( |
| 64 | self, root, split="train", break_mode="none", max_length=256, min_length=1 |
| 65 | ): |
| 66 | |
| 67 | if split not in ["train", "validation", "test"]: |
| 68 | raise ValueError( |
| 69 | 'Unknown split "{0}" for Wikitext-2. The split ' |
| 70 | "must be one of {train, validation, test}.".format(split) |
| 71 | ) |
| 72 | if break_mode not in ["none", "lines", "complete"]: |
| 73 | raise ValueError( |
| 74 | 'Unknown break_mode "{0}" for Wikitext-2. The ' |
| 75 | "break_mode must be one of {none, lines, complete}.".format(break_mode) |
| 76 | ) |
| 77 | |
| 78 | self.root = os.path.expanduser(root) |
| 79 | self.split = split |
| 80 | self.break_mode = break_mode |
| 81 | self.max_length = max_length |
| 82 | self.min_length = min_length |
| 83 | |
| 84 | # Load the dataset |
| 85 | filename = os.path.join(self.root, "wiki.{0}.npz".format(split)) |
| 86 | with open(filename, "rb") as f: |
| 87 | data = np.load(f) |
| 88 | self._tokens = torch.from_numpy(data["tokens"].astype(np.int64)) |
| 89 | self._sizes = tuple(data["sizes"]) |
| 90 | |
| 91 | # Create the examples depending on the break mode |
| 92 | if self.break_mode == "none": |
| 93 | indices = self._get_split_indices(self.num_tokens) |
| 94 | elif self.break_mode == "lines": |
| 95 | indices = self._get_split_indices_lines() |
| 96 | elif self.break_mode == "complete": |
| 97 | indices = self._get_split_indices_complete() |
| 98 | |
| 99 | self._examples = [ |
| 100 | self._tokens[start : start + length] |
| 101 | for (start, length) in indices |
| 102 | if length > self.min_length |
| 103 | ] |
| 104 | |
| 105 | def __getitem__(self, index): |
| 106 | example = self._examples[index] |
nothing calls this directly
no test coverage detected