[feat] merge

Author: tswsxk

EduNLP/Pretrain/gensim_vec.py

@@ -7,31 +7,60 @@
 from gensim.models.doc2vec import TaggedDocument
 from gensim.models.callbacks import CallbackAny2Vec
 from EduNLP.SIF.sif import sif4sci
-from EduNLP.Vector import D2V
+from EduNLP.Vector import D2V, BowLoader
 from copy import deepcopy
 import itertools as it
 
 __all__ = ["GensimWordTokenizer", "train_vector", "GensimSegTokenizer"]
 
 
 class GensimWordTokenizer(object):
-    def __init__(self, symbol="gm"):
+    def __init__(self, symbol="gm", general=False):
         """
 
         Parameters
         ----------
         symbol:
             gm
             fgm
+            gmas
+            fgmas
+        general:
+            True when item isn't in standard format, and want to tokenize formulas(except formulas in figure) linearly.
+            False when use 'ast' mothed to tokenize formulas instead of 'linear'.
+
+        Returns
+        ----------
+
+        Examples
+        ----------
+        >>> tokenizer = GensimWordTokenizer(symbol="gmas", general=True)
+        >>> token_item = tokenizer("有公式$\\FormFigureID{wrong1?}$，如图$\\FigureID{088f15ea-xxx}$,\
+        ... 若$x,y$满足约束条件公式$\\FormFigureBase64{wrong2?}$,$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$")
+        >>> print(token_item.tokens[:10])
+        ['公式', '[FORMULA]', '如图', '[FIGURE]', 'x', ',', 'y', '约束条件', '公式', '[FORMULA]']
+        >>> tokenizer = GensimWordTokenizer(symbol="fgmas", general=False)
+        >>> token_item = tokenizer("有公式$\\FormFigureID{wrong1?}$，如图$\\FigureID{088f15ea-xxx}$,\
+        ... 若$x,y$满足约束条件公式$\\FormFigureBase64{wrong2?}$,$\\SIFSep$，则$z=x+7 y$的最大值为$\\SIFBlank$")
+        >>> print(token_item.tokens[:10])
+        ['公式', '[FORMULA]', '如图', '[FIGURE]', '[FORMULA]', '约束条件', '公式', '[FORMULA]', '[SEP]', '[FORMULA]']
         """
         self.symbol = symbol
-        self.tokenization_params = {
-            "formula_params": {
-                "method": "ast",
-                "return_type": "list",
-                "ord2token": True
+        if general is True:
+            self.tokenization_params = {
+                "formula_params": {
+                    "method": "linear",
+                    "symbolize_figure_formula": True
+                }
+            }
+        else:
+            self.tokenization_params = {
+                "formula_params": {
+                    "method": "ast",
+                    "return_type": "list",
+                    "ord2token": True
+                }
             }
-        }
 
     def batch_process(self, *items):
         pass
@@ -135,16 +164,16 @@ def train_vector(items, w2v_prefix, embedding_dim=None, method="sg", binary=None
         binary = binary if binary is not None else True
     elif method == "tfidf":
         dictionary_path = train_vector(items, w2v_prefix, method="bow")
-        dictionary = D2V(dictionary_path, method="bow")
-        corpus = [dictionary(item) for item in items]
+        dictionary = BowLoader(dictionary_path)
+        corpus = [dictionary.infer_vector(item) for item in items]
         model = gensim.models.TfidfModel(corpus)
         binary = binary if binary is not None else True
     else:
         raise ValueError("Unknown method: %s" % method)
 
     filepath = w2v_prefix + method
     if embedding_dim is not None:
-        filepath = w2v_prefix + "_" + str(embedding_dim)
+        filepath = filepath + "_" + str(embedding_dim)
 
     if binary is True:
         filepath += ".bin"

EduNLP/SIF/tokenization/tokenization.py

@@ -33,6 +33,10 @@ def __init__(self, segment_list: SegmentList, text_params=None, formula_params=N
             "s": []
         }
         self.text_params = text_params if text_params is not None else {}
+        if formula_params is not None and "symbolize_figure_formula" in formula_params:
+            self.symbolize_figure_formula = formula_params.pop("symbolize_figure_formula")
+        else:
+            self.symbolize_figure_formula = False
         self.formula_params = formula_params if formula_params is not None else {"method": "linear"}
         self.formula_tokenize_method = self.formula_params.get("method")
         self.figure_params = figure_params if figure_params is not None else {}
@@ -166,6 +170,9 @@ def append_formula(self, segment, symbol=False, init=True):
             if symbol is True:
                 self._formula_tokens.append(len(self._tokens))
                 self._tokens.append(segment)
+            elif self.symbolize_figure_formula and isinstance(segment, FigureFormulaSegment):
+                self._formula_tokens.append(len(self._tokens))
+                self._tokens.append(Symbol(FORMULA_SYMBOL))
             elif isinstance(segment, FigureFormulaSegment):
                 self._formula_tokens.append(len(self._tokens))
                 self._tokens.append(segment)

EduNLP/Vector/__init__.py

@@ -1,7 +1,7 @@
 # coding: utf-8
 # 2021/5/29 @ tongshiwei
 
-from .gensim_vec import W2V, D2V
+from .gensim_vec import W2V, D2V, BowLoader, TfidfLoader
 from .const import *
 from .rnn import RNNModel
 from .t2v import T2V, get_pretrained_t2v

EduNLP/Vector/gensim_vec.py

@@ -11,7 +11,7 @@
 
 
 class W2V(Vector):
-    def __init__(self, filepath, method, binary=None):
+    def __init__(self, filepath, method=None, binary=None):
         """
 
         Parameters
@@ -69,16 +69,41 @@ def infer_tokens(self, items, *args, **kwargs) -> list:
         return [list(self(*item)) for item in items]
 
 
+class BowLoader(object):
+    def __init__(self, filepath):
+        self.dictionary = corpora.Dictionary.load(filepath)
+
+    def infer_vector(self, item, return_vec=False):
+        item = self.dictionary.doc2bow(item)
+        if not return_vec:
+            return item  # return dic as default
+        vec = [0 for i in range(len(self.dictionary.keys()))]
+        for i, v in item:
+            vec[i] = v
+        return vec
+
+    @property
+    def vector_size(self):
+        return len(self.dictionary.keys())
+
+
 class TfidfLoader(object):
     def __init__(self, filepath):
         self.tfidf_model = TfidfModel.load(filepath)
         # 'tfidf' model shold be used based on 'bow' model
         dictionary_path = re.sub(r"(.*)tfidf", r"\1bow", filepath)
         self.dictionary = corpora.Dictionary.load(dictionary_path)
 
-    def infer_vector(self, item):
-        item = self.dictionary.doc2bow(item)
-        return self.tfidf_model[item]
+    def infer_vector(self, item, return_vec=False):
+        dic_item = self.dictionary.doc2bow(item)
+        tfidf_item = self.tfidf_model[dic_item]
+        # return dic as default
+        if not return_vec:
+            return tfidf_item  # pragma: no cover
+        vec = [0 for i in range(len(self.dictionary.keys()))]
+        for i, v in tfidf_item:
+            vec[i] = v
+        return vec
 
     @property
     def vector_size(self):
@@ -92,24 +117,26 @@ def __init__(self, filepath, method="d2v"):
         if self._method == "d2v":
             self.d2v = Doc2Vec.load(filepath)
         elif self._method == "bow":
-            self.d2v = corpora.Dictionary.load(filepath)
+            self.d2v = BowLoader(filepath)
         elif self._method == "tfidf":
             self.d2v = TfidfLoader(filepath)
         else:
             raise ValueError("Unknown method: %s" % method)
 
     def __call__(self, item):
-        if self._method == "bow":
-            return self.d2v.doc2bow(item)
-        else:
+        if self._method == "d2v":
             return self.d2v.infer_vector(item)
+        else:
+            return self.d2v.infer_vector(item, return_vec=True)
 
     @property
     def vector_size(self):
-        if self._method in {"d2v", "tfidf"}:
+        if self._method == "d2v":
             return self.d2v.vector_size
         elif self._method == "bow":
-            return len(self.d2v.token2id)
+            return self.d2v.vector_size
+        elif self._method == "tfidf":
+            return self.d2v.vector_size
 
     def infer_vector(self, items, *args, **kwargs) -> list:
         return [self(item) for item in items]