chore: import upstream snapshot with attribution

python/paddle/text/datasets/imikolov.py

@@ -0,0 +1,204 @@
+#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import collections
+import tarfile
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+
+from paddle.dataset.common import _check_exists_and_download
+from paddle.io import Dataset
+
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+    _ImikolovDataType = Literal["NGRAM", "SEQ"]
+    _ImikolovDataSetMode = Literal["train", "test"]
+__all__ = []
+
+URL = 'https://dataset.bj.bcebos.com/imikolov%2Fsimple-examples.tgz'
+MD5 = '30177ea32e27c525793142b6bf2c8e2d'
+
+
+class Imikolov(Dataset):
+    """
+    Implementation of imikolov dataset.
+
+    Args:
+        data_file(str|None): path to data tar file, can be set None if
+            :attr:`download` is True. Default None.
+        data_type(str): 'NGRAM' or 'SEQ'. Default 'NGRAM'.
+        window_size(int): sliding window size for 'NGRAM' data. Default -1.
+        mode(str): 'train' 'test' mode. Default 'train'.
+        min_word_freq(int): minimal word frequencies for building word dictionary. Default 50.
+        download(bool): whether to download dataset automatically if
+            :attr:`data_file` is not set. Default True
+
+    Returns:
+        Dataset: instance of imikolov dataset
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> # doctest: +TIMEOUT(60)
+            >>> import paddle
+            >>> from paddle.text.datasets import Imikolov
+
+            >>> class SimpleNet(paddle.nn.Layer):
+            ...     def __init__(self):
+            ...         super().__init__()
+            ...
+            ...     def forward(self, src, trg):
+            ...         return paddle.sum(src), paddle.sum(trg)
+
+
+            >>> imikolov = Imikolov(mode='train', data_type='SEQ', window_size=2)
+
+            >>> for i in range(10):
+            ...     src, trg = imikolov[i]
+            ...     src = paddle.to_tensor(src)
+            ...     trg = paddle.to_tensor(trg)
+            ...
+            ...     model = SimpleNet()
+            ...     src, trg = model(src, trg)
+            ...     print(src.item(), trg.item())
+            2076 2075
+            2076 2075
+            675 674
+            4 3
+            464 463
+            2076 2075
+            865 864
+            2076 2075
+            2076 2075
+            1793 1792
+
+    """
+
+    data_file: str | None
+    data_type: _ImikolovDataType
+    window_size: int
+    mode: _ImikolovDataSetMode
+    min_word_freq: int
+    word_idx: dict[str, int]
+
+    def __init__(
+        self,
+        data_file: str | None = None,
+        data_type: _ImikolovDataType = 'NGRAM',
+        window_size: int = -1,
+        mode: _ImikolovDataSetMode = 'train',
+        min_word_freq: int = 50,
+        download: bool = True,
+    ) -> None:
+        assert data_type.upper() in [
+            'NGRAM',
+            'SEQ',
+        ], f"data type should be 'NGRAM', 'SEQ', but got {data_type}"
+        self.data_type = data_type.upper()
+
+        assert mode.lower() in [
+            'train',
+            'test',
+        ], f"mode should be 'train', 'test', but got {mode}"
+        self.mode = mode.lower()
+
+        self.window_size = window_size
+        self.min_word_freq = min_word_freq
+
+        self.data_file = data_file
+        if self.data_file is None:
+            assert download, (
+                "data_file is not set and downloading automatically disabled"
+            )
+            self.data_file = _check_exists_and_download(
+                data_file, URL, MD5, 'imikolov', download
+            )
+
+        # Build a word dictionary from the corpus
+        self.word_idx = self._build_work_dict(min_word_freq)
+
+        # read dataset into memory
+        self._load_anno()
+
+    def word_count(self, f, word_freq=None):
+        if word_freq is None:
+            word_freq = collections.defaultdict(int)
+
+        for l in f:
+            for w in l.strip().split():
+                word_freq[w] += 1
+            word_freq['<s>'] += 1
+            word_freq['<e>'] += 1
+
+        return word_freq
+
+    def _build_work_dict(self, cutoff: int) -> dict[str, int]:
+        train_filename = './simple-examples/data/ptb.train.txt'
+        test_filename = './simple-examples/data/ptb.valid.txt'
+        with tarfile.open(self.data_file) as tf:
+            trainf = tf.extractfile(train_filename)
+            testf = tf.extractfile(test_filename)
+            word_freq = self.word_count(testf, self.word_count(trainf))
+            if '<unk>' in word_freq:
+                # remove <unk> for now, since we will set it as last index
+                del word_freq['<unk>']
+
+            word_freq = [
+                x for x in word_freq.items() if x[1] > self.min_word_freq
+            ]
+
+            word_freq_sorted = sorted(word_freq, key=lambda x: (-x[1], x[0]))
+            words, _ = list(zip(*word_freq_sorted))
+            word_idx = dict(list(zip(words, range(len(words)))))
+            word_idx['<unk>'] = len(words)
+
+        return word_idx
+
+    def _load_anno(self) -> None:
+        self.data = []
+        with tarfile.open(self.data_file) as tf:
+            filename = f'./simple-examples/data/ptb.{self.mode}.txt'
+            f = tf.extractfile(filename)
+
+            UNK = self.word_idx['<unk>']
+            for l in f:
+                if self.data_type == 'NGRAM':
+                    assert self.window_size > -1, 'Invalid gram length'
+                    l = ["<s>", *l.strip().split(), "<e>"]
+                    if len(l) >= self.window_size:
+                        l = [self.word_idx.get(w, UNK) for w in l]
+                        for i in range(self.window_size, len(l) + 1):
+                            self.data.append(tuple(l[i - self.window_size : i]))
+                elif self.data_type == 'SEQ':
+                    l = l.strip().split()
+                    l = [self.word_idx.get(w, UNK) for w in l]
+                    src_seq = [self.word_idx["<s>"], *l]
+                    trg_seq = [*l, self.word_idx["<e>"]]
+                    if self.window_size > 0 and len(src_seq) > self.window_size:
+                        continue
+                    self.data.append((src_seq, trg_seq))
+                else:
+                    raise AssertionError('Unknown data type')
+
+    def __getitem__(
+        self, idx: int
+    ) -> tuple[npt.NDArray[np.int_], npt.NDArray[np.int_]]:
+        return tuple([np.array(d) for d in self.data[idx]])
+
+    def __len__(self) -> int:
+        return len(self.data)