Merge pull request #22 from bigdata-ustc/i2v

[FEATURE] add dynamic embedding model familty of rnn

Author: tswsxk

.travis.yml

@@ -4,6 +4,7 @@
 import torch
 from torch import nn
 from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
+from baize.torch import load_net
 
 
 class LM(nn.Module):
@@ -28,7 +29,7 @@ class LM(nn.Module):
     """
 
     def __init__(self, rnn_type: str, vocab_size: int, embedding_dim: int, hidden_size: int, num_layers=1,
-                 bidirectional=False, embedding=None, **kwargs):
+                 bidirectional=False, embedding=None, model_params=None, **kwargs):
         super(LM, self).__init__()
         rnn_type = rnn_type.upper()
         self.embedding = torch.nn.Embedding(vocab_size, embedding_dim) if embedding is None else embedding
@@ -61,12 +62,15 @@ def __init__(self, rnn_type: str, vocab_size: int, embedding_dim: int, hidden_si
             self.num_layers *= 2
         self.hidden_size = hidden_size
 
+        if model_params:
+            load_net(model_params, self, allow_missing=True)
+
     def forward(self, seq_idx, seq_len):
         seq = self.embedding(seq_idx)
         pack = pack_padded_sequence(seq, seq_len, batch_first=True)
-        h0 = torch.randn(self.num_layers, seq.shape[0], self.hidden_size)
+        h0 = torch.zeros(self.num_layers, seq.shape[0], self.hidden_size)
         if self.c is True:
-            c0 = torch.randn(self.num_layers, seq.shape[0], self.hidden_size)
+            c0 = torch.zeros(self.num_layers, seq.shape[0], self.hidden_size)
             output, (hn, _) = self.rnn(pack, (h0, c0))
         else:
             output, hn = self.rnn(pack, h0)

EduNLP/ModelZoo/rnn/rnn.py

@@ -3,3 +3,4 @@
 
 from .padder import PadSequence, pad_sequence
 from .device import set_device
+from .masker import Masker

EduNLP/ModelZoo/utils/__init__.py

@@ -0,0 +1,59 @@
+# coding: utf-8
+# 2021/8/3 @ tongshiwei
+
+from copy import deepcopy
+import numpy as np
+
+
+class Masker(object):
+    """
+    Examples
+    -------
+    >>> masker = Masker(per=0.5, seed=10)
+    >>> items = [[1, 1, 3, 4, 6], [2], [5, 9, 1, 4]]
+    >>> masked_seq, mask_label = masker(items)
+    >>> masked_seq
+    [[1, 1, 0, 0, 6], [2], [0, 9, 0, 4]]
+    >>> mask_label
+    [[0, 0, 1, 1, 0], [0], [1, 0, 1, 0]]
+    >>> items = [[1, 2, 3], [1, 1, 0], [2, 0, 0]]
+    >>> masked_seq, mask_label = masker(items, [3, 2, 1])
+    >>> masked_seq
+    [[1, 0, 3], [0, 1, 0], [2, 0, 0]]
+    >>> mask_label
+    [[0, 1, 0], [1, 0, 0], [0, 0, 0]]
+    >>> masker = Masker(mask="[MASK]", per=0.5, seed=10)
+    >>> items = [["a", "b", "c"], ["d", "[PAD]", "[PAD]"], ["hello", "world", "[PAD]"]]
+    >>> masked_seq, mask_label = masker(items, length=[3, 1, 2])
+    >>> masked_seq
+    [['a', '[MASK]', 'c'], ['d', '[PAD]', '[PAD]'], ['hello', '[MASK]', '[PAD]']]
+    >>> mask_label
+    [[0, 1, 0], [0, 0, 0], [0, 1, 0]]
+    """
+
+    def __init__(self, mask: (int, str, ...) = 0, per=0.2, seed=None):
+        """
+
+        Parameters
+        ----------
+        mask: int, str
+        per
+        seed
+        """
+        self.seed = np.random.default_rng(seed)
+        self.per = per
+        self.mask = mask
+
+    def __call__(self, seqs, length=None, *args, **kwargs) -> tuple:
+        seqs = deepcopy(seqs)
+        masked_list = []
+        if length is None:
+            length = [len(seq) for seq in seqs]
+        for seq, _length in zip(seqs, length):
+            masked = self.seed.choice(len(seq) - 1, size=int(_length * self.per), replace=False)
+            _masked_list = [0] * len(seq)
+            for _masked in masked:
+                seq[_masked] = self.mask
+                _masked_list[_masked] = 1
+            masked_list.append(_masked_list)
+        return seqs, masked_list

EduNLP/ModelZoo/utils/masker.py

@@ -19,7 +19,7 @@ def __init__(self, w2v: (W2V, tuple, list, dict, None), freeze=True, device=None
         elif isinstance(w2v, W2V):
             self.w2v = w2v
         else:
-            raise TypeError("w2v argument must be one of W2V, tuple, list, dict or None")
+            raise TypeError("w2v argument must be one of W2V, tuple, list, dict or None, now is %s" % type(w2v))
 
         if self.w2v is not None:
             self.vocab_size = len(self.w2v)
@@ -63,7 +63,10 @@ def indexing(self, items: List[List[str]], padding=False, indexing=True) -> tupl
 
         Returns
         -------
-        word_id: list of list of int
+        token_idx: list of list of int
+            the list of the tokens of each item
+        token_len: list of int
+            the list of the length of tokens of each item
         """
         items_idx = [[self.key_to_index(word) for word in item] for item in items] if indexing else items
         item_len = [len(_idx) for _idx in items_idx]

EduNLP/Vector/embedding.py

@@ -11,3 +11,10 @@ def infer_tokens(self, items, *args, **kwargs) -> ...:
     @property
     def vector_size(self):
         raise NotImplementedError
+
+    @property
+    def is_frozen(self):  # pragma: no cover
+        return True
+
+    def freeze(self, *args, **kwargs):  # pragma: no cover
+        pass

EduNLP/Vector/meta.py

@@ -6,6 +6,7 @@
 from ..embedding import Embedding
 from ..meta import Vector
 from EduNLP.ModelZoo import rnn, set_device
+from baize.torch import save_params
 
 
 class RNNModel(Vector):
@@ -38,7 +39,8 @@ class RNNModel(Vector):
     torch.Size([2, 3, 2])
     """
 
-    def __init__(self, rnn_type, w2v: (W2V, tuple, list, dict, None), hidden_size, freeze_pretrained=True, device=None,
+    def __init__(self, rnn_type, w2v: (W2V, tuple, list, dict, None), hidden_size,
+                 freeze_pretrained=True, model_params=None, device=None,
                  **kwargs):
         self.embedding = Embedding(w2v, freeze_pretrained, **kwargs)
         for key in ["vocab_size", "embedding_dim"]:
@@ -50,6 +52,7 @@ def __init__(self, rnn_type, w2v: (W2V, tuple, list, dict, None), hidden_size, f
             self.embedding.embedding_dim,
             hidden_size=hidden_size,
             embedding=self.embedding.embedding,
+            model_params=model_params,
             **kwargs
         )
         self.bidirectional = self.rnn.rnn.bidirectional
@@ -86,3 +89,22 @@ def vector_size(self) -> int:
 
     def set_device(self, device):
         self.rnn = set_device(self.rnn, device)
+
+    def save(self, filepath, save_embedding=False):
+        save_params(filepath, self.rnn, select=None if save_embedding is True else '^(?!.*embedding)')
+        return filepath
+
+    def freeze(self, *args, **kwargs):
+        return self.eval()
+
+    @property
+    def is_frozen(self):
+        return not self.rnn.training
+
+    def eval(self):
+        self.rnn.eval()
+        return self
+
+    def train(self, mode=True):
+        self.rnn.train(mode)
+        return self

EduNLP/Vector/rnn/rnn.py

@@ -0,0 +1,22 @@
+# coding: utf-8
+# 2021/8/3 @ tongshiwei
+
+from longling import load_jsonl
+from EduNLP.Tokenizer import get_tokenizer
+from EduNLP.Pretrain import train_vector
+from EduNLP.Vector import W2V, RNNModel
+
+
+def etl():
+    tokenizer = get_tokenizer("text")
+    return tokenizer([item["stem"] for item in load_jsonl("../../../data/OpenLUNA.json")])
+
+
+items = list(etl())
+model_path = train_vector(items, "./w2v", 10, "sg")
+
+w2v = W2V(model_path, "sg")
+rnn = RNNModel("lstm", w2v, 5, device="cpu")
+saved_params = rnn.save("./lstm.params", save_embedding=True)
+
+rnn1 = RNNModel("lstm", w2v, 5, model_params=saved_params)

examples/pretrain/rnn/rnn.py

@@ -21,7 +21,7 @@
 
 setup(
     name='EduNLP',
-    version='0.0.3',
+    version='0.0.4',
     extras_require={
         'test': test_deps,
         'tutor': tutor_deps,
@@ -35,11 +35,23 @@
         'jieba',
         'js2py',
         'torch',
-        'EduData>=0.0.16'
+        'EduData>=0.0.16',
+        'PyBaize[torch]>=0.0.3'
     ],  # And any other dependencies foo needs
     entry_points={
         "console_scripts": [
             "edunlp = EduNLP.main:cli",
         ],
     },
+    classifiers=[
+        'Programming Language :: Python :: 3.6',
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        "Environment :: Other Environment",
+        "Intended Audience :: Developers",
+        "License :: OSI Approved :: Apache License 2.0 (Apache 2.0)",
+        "Operating System :: OS Independent",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
 )

setup.py

@@ -1,6 +1,7 @@
 # coding: utf-8
 # 2021/5/30 @ tongshiwei
 
+import torch
 import numpy as np
 import pytest
 from EduNLP.Pretrain import train_vector, GensimWordTokenizer
@@ -111,10 +112,22 @@ def test_rnn(stem_tokens, tmpdir):
         item = rnn.infer_vector(stem_tokens[:1])
         assert tokens.shape == (1, len(stem_tokens[0]), 20 * (2 if rnn.bidirectional else 1))
         assert item.shape == (1, rnn.vector_size)
+        item_vec = rnn.infer_vector(stem_tokens[:1])
+        assert torch.equal(item, item_vec)
 
         t2v = T2V(rnn_type, w2v, 20)
         assert len(t2v(stem_tokens[:1])[0]) == t2v.vector_size
 
+        saved_params = rnn.save(str((tmpdir / method).join("stem_tf_rnn.params")), save_embedding=True)
+
+        rnn = RNNModel(rnn_type, w2v, 20, device="cpu", model_params=saved_params)
+        rnn.train()
+        assert rnn.is_frozen is False
+        rnn.freeze()
+        assert rnn.is_frozen is True
+        item_vec1 = rnn.infer_vector(stem_tokens[:1])
+        assert torch.equal(item, item_vec1)
+
 
 def test_d2v(stem_tokens, tmpdir, stem_data):
     method = "d2v"