From be5595fc5f40f7d281f9318ff26095c0d15ed5da Mon Sep 17 00:00:00 2001
From: LYH-YF <yifan2250@gmail.com>
Date: Thu, 14 Oct 2021 22:22:48 +0800
Subject: [PATCH] mwptoolkit v0.0.5

---
 .gitignore                                    |    8 +
 PYPI.md                                       |    1 +
 README.md                                     |    2 +
 docs/_static/cmd.html                         |    5 +-
 docs/conf.py                                  |    2 +-
 mwptoolkit/config/__init__.py                 |    3 +-
 mwptoolkit/data/__init__.py                   |    3 +-
 mwptoolkit/data/dataloader/__init__.py        |   10 +-
 .../dataloader/multi_equation_dataloader.py   |    9 +-
 .../data/dataloader/pretrain_dataloader.py    |   12 +-
 .../dataloader/single_equation_dataloader.py  |    9 +-
 mwptoolkit/data/dataset/__init__.py           |   10 +-
 mwptoolkit/evaluate/__init__.py               |    4 +-
 mwptoolkit/evaluate/equ_solver.py             |  241 --
 mwptoolkit/evaluate/evaluator.py              |  104 +-
 mwptoolkit/loss/__init__.py                   |   10 +-
 mwptoolkit/model/Graph2Tree/__init__.py       |    5 +-
 mwptoolkit/model/PreTrain/__init__.py         |    6 +-
 mwptoolkit/model/Seq2Seq/__init__.py          |   12 +-
 mwptoolkit/model/Seq2Seq/mathen.py            |    2 +-
 mwptoolkit/model/Seq2Tree/__init__.py         |    8 +-
 mwptoolkit/model/Seq2Tree/sausolver.py        |  756 ++++-
 mwptoolkit/model/__init__.py                  |    2 +-
 mwptoolkit/module/Decoder/tree_decoder.py     |  152 +-
 mwptoolkit/module/Layer/tree_layers.py        |   36 +-
 mwptoolkit/trainer/supervised_trainer.py      |  269 +-
 mwptoolkit/trainer/template_trainer.py        |    3 +-
 mwptoolkit/trainer/trainer.py                 | 2739 -----------------
 setup.py                                      |    2 +-
 29 files changed, 1289 insertions(+), 3136 deletions(-)
 delete mode 100644 mwptoolkit/evaluate/equ_solver.py
 delete mode 100644 mwptoolkit/trainer/trainer.py

diff --git a/.gitignore b/.gitignore
index 174d1e2..eb3c63a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,7 @@
 /.vscode
+/.idea
+/mwptoolkit.egg-info
+/dist
 __pycache__
 /dataset/math23k/partition.py
 /dataset/ape200k/process.py
@@ -6,12 +9,17 @@ __pycache__
 /dataset/mawps/MAWPS_.json
 /dataset/hmwp/process.py
 /dataset/math23k/graph2tree_deprel_info.json
+/dataset/alg514/stat_equations.py
 preprocess.py
 pos_info.json
 pos_info_.json
 span_level_deprel_tree_info.json
 deprel_tree_info.json
 *.pth
+.pypirc
+make.bat
+Makefile
+run_setup.sh
 /pretrain/gpt2_cn
 /pretrain/gpt2_en
 *.log
diff --git a/PYPI.md b/PYPI.md
index c5d3301..80986c9 100644
--- a/PYPI.md
+++ b/PYPI.md
@@ -8,4 +8,5 @@ MWPToolkit is a PyTorch-based toolkit for Math Word Problem(MWP) solving. It is
 * **Highly modularized framework**. MWP toolkit is designed with highly reused modules and provides convenient interfaces for users. Specifically, data preprocessor, data loader, encoder, decoder and evaluator form the running procedure. Each module could be developed and extended independently.
 
 ## News
+* **Fix some bugs.**
 
diff --git a/README.md b/README.md
index 49311d8..e8ae563 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,8 @@ Our framework has the following architecture. You could utilize our toolkit to e
 
 ## News
 
+* **Fix some bugs.**
+
 ## Characteristics
 
 * **Unification and Modularization**. We decouple solvers with different model architectures into highly modularized, reusable components and integrate them in a unified framework, which includes data, model, evaluation modules. It is convenient for you to study MWPs at a conceptual level and compare different models fairly.
diff --git a/docs/_static/cmd.html b/docs/_static/cmd.html
index 9a25c1c..dffea40 100644
--- a/docs/_static/cmd.html
+++ b/docs/_static/cmd.html
@@ -7,7 +7,7 @@
   
   <meta name="viewport" content="width=device-width, initial-scale=1.0" />
   
-  <title>mwptoolkit.config.configuration &mdash; MWPToolkit 0.0.4 documentation</title>
+  <title>command line &mdash; MWPToolkit 0.0.5 documentation</title>
   
 
   
@@ -73,7 +73,7 @@
             cmd_line = cmd_line+"--dataset="+dataset+" ";
 
             var single_dataset=new Array("math23k","asdiv-a","mawps-single","mawps_asdiv-a_svamp");
-            var multi_dataset=new Array("draw","mawps","hmwp");
+            var multi_dataset=new Array("draw","mawps","hmwp","alg514");
             if (single_dataset.includes(dataset)){
                 cmd_line = cmd_line+"--task_type="+"single_equation ";
             }
@@ -316,6 +316,7 @@ <h1>command line<a class="headerlink" href="#cmd.html" title="Permalink to this
                             <td>
                                 <input id="dataset" style="width: 100%; height: 100%" name="dataset" type="text" list="datasetlist" />
                                 <datalist id="datasetlist">
+                                    <option value="alg514" />
                                     <option value="asdiv-a" />
                                     <option value="draw" />
                                     <option value="hmwp" />
diff --git a/docs/conf.py b/docs/conf.py
index 57c2fbe..29774b0 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -23,7 +23,7 @@
 author = "''"
 
 # The full version, including alpha/beta/rc tags
-release = '0.0.4'
+release = '0.0.5'
 
 
 # -- General configuration ---------------------------------------------------
diff --git a/mwptoolkit/config/__init__.py b/mwptoolkit/config/__init__.py
index 44a8fa6..aec3b6d 100644
--- a/mwptoolkit/config/__init__.py
+++ b/mwptoolkit/config/__init__.py
@@ -1,4 +1,5 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.config import configuration
\ No newline at end of file
+from mwptoolkit.config import configuration
+from mwptoolkit.config.configuration import Config
diff --git a/mwptoolkit/data/__init__.py b/mwptoolkit/data/__init__.py
index 43eb5a7..9608560 100644
--- a/mwptoolkit/data/__init__.py
+++ b/mwptoolkit/data/__init__.py
@@ -1 +1,2 @@
-from mwptoolkit.data import dataloader,dataset,utils
\ No newline at end of file
+from mwptoolkit.data import dataloader,dataset,utils
+
diff --git a/mwptoolkit/data/dataloader/__init__.py b/mwptoolkit/data/dataloader/__init__.py
index c25f700..ec7d762 100644
--- a/mwptoolkit/data/dataloader/__init__.py
+++ b/mwptoolkit/data/dataloader/__init__.py
@@ -1,4 +1,12 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.data.dataloader import abstract_dataloader,dataloader_ept,dataloader_multiencdec,pretrain_dataloader,single_equation_dataloader,multi_equation_dataloader,template_dataloader
\ No newline at end of file
+from mwptoolkit.data.dataloader import abstract_dataloader,dataloader_ept,dataloader_multiencdec,pretrain_dataloader,single_equation_dataloader,multi_equation_dataloader,template_dataloader
+
+from mwptoolkit.data.dataloader.abstract_dataloader import AbstractDataLoader
+from mwptoolkit.data.dataloader.single_equation_dataloader import SingleEquationDataLoader
+from mwptoolkit.data.dataloader.multi_equation_dataloader import MultiEquationDataLoader
+from mwptoolkit.data.dataloader.template_dataloader import TemplateDataLoader
+from mwptoolkit.data.dataloader.dataloader_ept import DataLoaderEPT
+from mwptoolkit.data.dataloader.dataloader_multiencdec import DataLoaderMultiEncDec
+from mwptoolkit.data.dataloader.pretrain_dataloader import PretrainDataLoader
diff --git a/mwptoolkit/data/dataloader/multi_equation_dataloader.py b/mwptoolkit/data/dataloader/multi_equation_dataloader.py
index 2998c95..c929a43 100644
--- a/mwptoolkit/data/dataloader/multi_equation_dataloader.py
+++ b/mwptoolkit/data/dataloader/multi_equation_dataloader.py
@@ -58,13 +58,8 @@ def __init__(self, config, dataset):
         self.validset_nums = len(dataset.validset)
         self.testset_nums = len(dataset.testset)
 
-        if config["pretrained_model_path"]:
-            self.pretrained_tokenzier = AutoTokenizer.from_pretrained(config["pretrained_model_path"])
-            if config["model"].lower() in ["ept"]:
-                self.pretrained_tokenzier.add_special_tokens({'additional_special_tokens': ['[N]']})
-        else:
-            self.in_pad_token = dataset.in_word2idx[SpecialTokens.PAD_TOKEN]
-            self.in_unk_token = dataset.in_word2idx[SpecialTokens.UNK_TOKEN]
+        self.in_pad_token = dataset.in_word2idx[SpecialTokens.PAD_TOKEN]
+        self.in_unk_token = dataset.in_word2idx[SpecialTokens.UNK_TOKEN]
 
         if self.symbol_for_tree or self.equation_fix == FixType.MultiWayTree:
             self.out_pad_token = self.in_pad_token
diff --git a/mwptoolkit/data/dataloader/pretrain_dataloader.py b/mwptoolkit/data/dataloader/pretrain_dataloader.py
index bb02b2c..5a1d27f 100644
--- a/mwptoolkit/data/dataloader/pretrain_dataloader.py
+++ b/mwptoolkit/data/dataloader/pretrain_dataloader.py
@@ -142,14 +142,15 @@ def load_batch(self,batch_data):
         num_stack_batch = []
 
         group_nums_batch = []
-        for data in batch_data:
-            data['question_']=self.dataset.tokenizer.tokenize(' '.join(data["question"]))
-        batch_data=sorted(batch_data,key=lambda x:len(x['question_']),reverse=True)
+        # for data in batch_data:
+        #     data['question_']=self.dataset.tokenizer.tokenize(' '.join(data["question"]))
+        #batch_data=sorted(batch_data,key=lambda x:len(x['question_']),reverse=True)
+        batch_data = sorted(batch_data, key=lambda x: len(x['question']), reverse=True)
         for data in batch_data:
             ques_tensor = []
             equ_tensor = []
             temp_tensor = []
-            sentence = data["question_"]
+            sentence = data["question"]
             equation = data["equation"]
             template = data["template"]
 
@@ -281,7 +282,8 @@ def load_batch(self,batch_data):
     def _word2idx(self, sentence):
         sentence_idx = []
         
-        sentence_idx = self.dataset.tokenizer.convert_tokens_to_ids(sentence)
+        # sentence_idx = self.dataset.tokenizer.convert_tokens_to_ids(sentence)
+        sentence_idx = self.dataset.tokenizer.encode(sentence,add_special_token=False)
         
         return sentence_idx
     
diff --git a/mwptoolkit/data/dataloader/single_equation_dataloader.py b/mwptoolkit/data/dataloader/single_equation_dataloader.py
index 8573bc9..94d478c 100644
--- a/mwptoolkit/data/dataloader/single_equation_dataloader.py
+++ b/mwptoolkit/data/dataloader/single_equation_dataloader.py
@@ -56,13 +56,8 @@ def __init__(self, config, dataset):
         self.validset_nums = len(dataset.validset)
         self.testset_nums = len(dataset.testset)
 
-        if config["pretrained_model"]:
-            self.pretrained_tokenzier = AutoTokenizer.from_pretrained(config["pretrained_model"])
-            if config["model"].lower() in ["ept"]:
-                self.pretrained_tokenzier.add_special_tokens({'additional_special_tokens': ['[N]']})
-        else:
-            self.in_pad_token = dataset.in_word2idx[SpecialTokens.PAD_TOKEN]
-            self.in_unk_token = dataset.in_word2idx[SpecialTokens.UNK_TOKEN]
+        self.in_pad_token = dataset.in_word2idx[SpecialTokens.PAD_TOKEN]
+        self.in_unk_token = dataset.in_word2idx[SpecialTokens.UNK_TOKEN]
 
         if self.symbol_for_tree or self.equation_fix == FixType.MultiWayTree:
             self.out_pad_token = self.in_pad_token
diff --git a/mwptoolkit/data/dataset/__init__.py b/mwptoolkit/data/dataset/__init__.py
index 82bab40..3270531 100644
--- a/mwptoolkit/data/dataset/__init__.py
+++ b/mwptoolkit/data/dataset/__init__.py
@@ -1,4 +1,12 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.data.dataset import abstract_dataset,dataset_ept,dataset_multiencdec,pretrain_dataset,single_equation_dataset,multi_equation_dataset,template_dataset
\ No newline at end of file
+from mwptoolkit.data.dataset import abstract_dataset,dataset_ept,dataset_multiencdec,pretrain_dataset,single_equation_dataset,multi_equation_dataset,template_dataset
+
+from mwptoolkit.data.dataset.abstract_dataset import AbstractDataset
+from mwptoolkit.data.dataset.single_equation_dataset import SingleEquationDataset
+from mwptoolkit.data.dataset.multi_equation_dataset import MultiEquationDataset
+from mwptoolkit.data.dataset.template_dataset import TemplateDataset
+from mwptoolkit.data.dataset.dataset_ept import DatasetEPT
+from mwptoolkit.data.dataset.dataset_multiencdec import DatasetMultiEncDec
+from mwptoolkit.data.dataset.pretrain_dataset import PretrainDataset
diff --git a/mwptoolkit/evaluate/__init__.py b/mwptoolkit/evaluate/__init__.py
index c1d1ebd..341abf3 100644
--- a/mwptoolkit/evaluate/__init__.py
+++ b/mwptoolkit/evaluate/__init__.py
@@ -1,4 +1,6 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.evaluate import evaluator
\ No newline at end of file
+from mwptoolkit.evaluate import evaluator
+
+from mwptoolkit.evaluate.evaluator import PrefixEvaluator,PostfixEvaluator,InfixEvaluator,MultiWayTreeEvaluator,MultiEncDecEvaluator
\ No newline at end of file
diff --git a/mwptoolkit/evaluate/equ_solver.py b/mwptoolkit/evaluate/equ_solver.py
deleted file mode 100644
index cbac386..0000000
--- a/mwptoolkit/evaluate/equ_solver.py
+++ /dev/null
@@ -1,241 +0,0 @@
-def joint_number_(text_list): #match longer fraction such as ( 1 / 1000000 )
-    #text_list='( | ( 5 / 7 ) - ( 1 / 14 ) + ( 5 / 6 ) ) / ( 5 / 42 ) ．'.split(" ")
-    #text_list="计 ( 123 (1/10000) ) 算 ( 1 / 2 ) + ( 1 / 6 ) + ( 1 / 12 ) + ( 1 / 20 ) + … + ( 1 / 380 ) = 多 少 ．".split()
-    text_list="苹 果 树 比 梨 树 少 ( 3 / 8 ) ， 梨 树 比 苹 果 树 多 ( ( ( ) ) / ( ( ) ) )".split(" ")
-    new_list=[]
-    i=0
-    while i < len(text_list):
-        if text_list[i] == '(':
-            try:
-                j=text_list[i:].index(')')
-                if i+1==i+j:
-                    j=None
-                if "(" in text_list[i+1:i+j+1]:
-                    j=None
-            except:
-                j=None
-            if j:
-                stack=[]
-                flag=True
-                idx=0
-                for temp_idx,word in enumerate(text_list[i:i+j+1]):
-                    if word in ["(",")","/"] or word.isdigit():
-                        stack.append(word)
-                        idx=temp_idx
-                    else:
-                        flag=False
-                        break
-                if "/" not in stack:
-                    flag=False
-                if flag:
-                    number=''.join(stack)
-                    new_list.append(number)
-                else:
-                    for word in stack:
-                        new_list.append(word)
-                i+=idx+1
-            else:
-                new_list.append(text_list[i])
-                i+=1
-        else:
-            new_list.append(text_list[i])
-            i+=1
-    return new_list
-#joint_number_([])
-from typing import Counter
-import torch
-from torch import nn, stack
-from torch.nn import functional as F
-import stanza
-import copy
-def get_group_num(doc,sent_len,num_pos):
-    heads_deprel_upos=doc.get(['head','deprel','upos','text'])
-    num_pos=[pos-sent_len for pos in num_pos]
-    for idx,x in enumerate(heads_deprel_upos):
-        print(idx,x)
-    for n_pos in num_pos:
-        pos_stack=[]
-        group_num=[]
-        #group_num.append(n_pos)
-        pos_stack.append([n_pos,heads_deprel_upos[n_pos][1]])
-        count=0
-        head_pos=heads_deprel_upos[n_pos][0]
-        head_dep=heads_deprel_upos[n_pos][1]
-        for idx,x in enumerate(heads_deprel_upos):
-            if heads_deprel_upos[idx][0]==head_pos and n_pos!=idx:
-                deprel=heads_deprel_upos[idx][1]
-                pos_stack.append([idx,deprel])
-        while pos_stack:
-            count+=1
-            pos_dep=pos_stack.pop(0)
-            pos=pos_dep[0]
-            dep=pos_dep[1]
-            head_pos=heads_deprel_upos[pos][0]-1
-            upos=heads_deprel_upos[pos][2]
-            if upos not in ['NOUN','NUM','ADJ','VERB','DET', 'SYM']:
-                continue
-            elif upos == 'NOUN' and dep not in ['compound','nsubj:pass','nsubj','compound']:
-                continue
-            elif upos == 'VERB' and dep not in ['conj','root']:
-                continue
-            elif upos == 'ADJ' and dep not in ['amod']:
-                continue
-            elif upos == 'DET' and dep not in ['advmod']:
-                continue
-            elif upos == 'SYM' and dep not in ['obl']:
-                continue
-            else:
-                group_num.append(pos)
-            if head_pos>=0:
-                head_dep=heads_deprel_upos[head_pos][1]
-                if [head_pos,head_dep] in pos_stack:
-                    pass
-                else:
-                    pos_stack.append([head_pos,head_dep])
-        print(count)
-        print(group_num)
-    return []
-def get_group_num_(token_list,sent_len,num_pos):
-    group_nums=[]
-    num_pos=[pos-sent_len for pos in num_pos]
-    for n_pos in num_pos:
-        pos_stack=[]
-        group_num=[]
-        pos_stack.append([n_pos,token_list[n_pos]["deprel"]])
-        count=0
-        head_pos=token_list[n_pos]['head']
-        for idx,x in enumerate(token_list):
-            if x['head']==head_pos and n_pos!=idx:
-                deprel=x["deprel"]
-                pos_stack.append([idx,deprel])
-        while pos_stack:
-            count+=1
-            pos_dep=pos_stack.pop(0)
-            pos=pos_dep[0]
-            dep=pos_dep[1]
-            head_pos=token_list[pos]['head']-1
-            upos=token_list[pos]['upos']
-            if upos not in ['NOUN','NUM','ADJ','VERB','DET', 'SYM']:
-                continue
-            elif upos == 'NOUN' and dep not in ['compound','nsubj:pass','nsubj','compound']:
-                continue
-            elif upos == 'VERB' and dep not in ['conj','root']:
-                continue
-            elif upos == 'ADJ' and dep not in ['amod']:
-                continue
-            elif upos == 'DET' and dep not in ['advmod']:
-                continue
-            elif upos == 'SYM' and dep not in ['obl']:
-                continue
-            else:
-                group_num.append(pos+sent_len)
-            if head_pos>=0:
-                head_dep=token_list[head_pos]['deprel']
-                if [head_pos,head_dep] in pos_stack:
-                    pass
-                else:
-                    pos_stack.append([head_pos,head_dep])
-        group_nums.append(group_num)
-    return group_nums
-# nlp=stanza.Pipeline('en',processors='depparse,tokenize,pos,lemma',tokenize_pretokenized=True)
-# text="348 teddy bears are sold for $ 23 each . There are total 470 teddy bears in a store and the remaining teddy bears are sold for $ 17 each . How much did the store earn after selling all the teddy bears ."
-# doc=nlp(text)
-# print(doc)
-# num_pos=[0, 7,13,28]
-# token_lists=doc.to_dict()
-# sentences=text.split(".")
-# sent_len=[]
-# num_pos=[0, 7,13,28]
-# num_pos_list=[]
-# last_pos_list=[]
-# l=0
-# for idx,sentence in enumerate(sentences):
-#     sentence=sentence+' .'
-#     sentences[idx]=sentence
-#     sent_len.append(l)
-#     l+=len(sentence.split(" "))
-#     n_pos=[pos for pos in num_pos if pos <l ]
-#     for pos in n_pos:
-#         num_pos.remove(pos)
-#     num_pos_list.append(n_pos)
-# for idx,token_list in enumerate(token_lists):
-#     doc=nlp(sentences[idx])
-#     #print(doc)
-#     token_list=doc.to_dict()[0]
-#     group=get_group_num_(token_list,sent_len[idx],num_pos_list[idx])
-#     print(group)
-# exit()
-#text="在 一 正方形 花池 的 4 周 栽 了 44 棵 柳树 ， 每 两棵 柳树 之间 的 间隔 是 20 米 ， 这个 正方形 的 周长 = 多少 米 ？"
-#text="在一正方形花池的4周栽了44棵柳树，每两棵柳树之间的间隔是20米，这个正方形的周长=多少米？"
-#text=['在','一','正方形','花池','的','4','周','栽','了','44','棵','柳树','，','每','两棵', '柳树', '之间', '的', '间隔', '是', '20', '米', '，', '这个','正方形' ,'的', '周长', '=', '多少' ,'米', '？']
-# nlp=stanza.Pipeline('en',processors='depparse,tokenize,pos,lemma',tokenize_pretokenized=True)
-# text="348 teddy bears are sold for $ 23 each. There are total 470 teddy bears in a store and the remaining teddy bears are sold for $ 17 each. How much did the store earn after selling all the teddy bears."
-# sentences=text.split(".")
-# sent_len=[]
-# num_pos=[0, 7,13,28]
-# num_pos_list=[]
-# last_pos_list=[]
-# l=0
-# for idx,sentence in enumerate(sentences):
-#     sentence=sentence+' .'
-#     sentences[idx]=sentence
-#     sent_len.append(l)
-#     l+=len(sentence.split(" "))
-#     n_pos=[pos for pos in num_pos if pos <l ]
-#     for pos in n_pos:
-#         num_pos.remove(pos)
-#     num_pos_list.append(n_pos)
-# for sent,s_len,sub_n_pos in zip(sentences,sent_len,num_pos_list):
-#     doc=nlp(sent+' .')
-#     print(doc)
-#     g_n=get_group_num(doc,s_len,sub_n_pos)
-#doc = nlp("镇海 雅乐 学校 二年级 的 小朋友 到 一条 小路 的 一边 植树 。 小朋友 们 每隔 2 米 种 一棵树 （ 马路 两头 都 种 了 树 ） ， 最后 发现 一共 种 了 11 棵 ， 这 条 小路 长 多少 米 ．")
-#doc = zh_nlp(text)
-#doc.sentences[0].print_dependencies()
-#print(doc)
-#nlp = stanza.Pipeline('zh-hans') # This sets up a default neural pipeline in English
-#text='地图 与 现实 的 比例 是 1:500 。'
-#doc = nlp(text)
-#doc.sentences[0].print_dependencies()
-#doc = nlp("348 teddy bears are sold for $ 23 each. There are total 470 teddy bears in a store and the remaining teddy bears are sold for $ 17 each. How much did the store earn after selling all the teddy bears.")
-#print(doc)
-# doc.sentences[0].print_dependencies()
-# doc.sentences[1].print_dependencies()
-# doc.sentences[2].print_dependencies()
-print(1)
-# class SeqAttention(nn.Module):
-#     def __init__(self, hidden_size):
-#         super(SeqAttention, self).__init__()
-#         self.hidden_size=hidden_size
-
-#         self.linear_out = nn.Linear(hidden_size*2, hidden_size)
-
-#     def forward(self, output, encoder_outputs,mask):
-#         '''
-#         Args:
-#             hidden: shape [batch_size, 1, hidden_size].
-#             encoder_outputs: shape [batch_size, sequence_length, hidden_size].
-#         '''
-#         batch_size = output.size(0)
-#         seq_length = encoder_outputs.size(1)
-        
-#         attn = torch.bmm(output, encoder_outputs.transpose(1,2))
-#         if mask is not None:
-#             attn.data.masked_fill_(mask, -float('inf'))
-#         attn = F.softmax(attn.view(-1, seq_length), dim=1).view(batch_size, -1, seq_length)
-
-#         mix = torch.bmm(attn, encoder_outputs)
-
-#         combined = torch.cat((mix, output), dim=2)
-
-#         output = torch.tanh(self.linear_out(combined.view(-1, 2*self.hidden_size)))\
-#                             .view(batch_size, -1, self.hidden_size)
-
-#         # output: (b, o, dim)
-#         # attn  : (b, o, i)
-#         return output, attn
-# attn=SeqAttention(128)
-# x=torch.rand((4,20,128))
-# y=torch.rand((4,1,128))
-# output,atten=attn.forward(y,x,None)
-# print(output.size())
\ No newline at end of file
diff --git a/mwptoolkit/evaluate/evaluator.py b/mwptoolkit/evaluate/evaluator.py
index 0b947fd..8d5666a 100644
--- a/mwptoolkit/evaluate/evaluator.py
+++ b/mwptoolkit/evaluate/evaluator.py
@@ -197,7 +197,19 @@ def result_multi(self, test_exp, tar_exp):
             if flag == True:
                 return True, False, tar_exp, tar_exp
             else:
-                return False, False, tar_exp, tar_exp
+                flag = True
+                test_solves_list = list(test_solves.values())
+                target_solvers_list = list(tar_solves.values())
+                t1 = sorted(test_solves_list)
+                t2 = sorted(target_solvers_list)
+                for v1, v2 in zip(t1, t2):
+                    if abs(v1 - v2) > 1e-4:
+                        flag = False
+                        break
+                if flag:
+                    return True, False, test_exp, tar_exp
+                else:
+                    return False, False, test_exp, tar_exp
         except:
             return False, False, tar_exp, tar_exp
 
@@ -369,18 +381,18 @@ def result_multi(self, test_exp, tar_exp):
 
         Args:
             test_exp (list): list of test expression.
-            
+
             tar_exp (list): list of target expression.
 
         Returns:
             (tuple(bool,bool,list,list))
 
             val_ac (bool): the correctness of test expression answer compared to target expression answer.
-            
+
             equ_ac (bool): the correctness of test expression compared to target expression.
-            
+
             test_exp (list): list of test expression.
-            
+
             tar_exp (list): list of target expression.
         """
         if test_exp is []:
@@ -431,7 +443,19 @@ def result_multi(self, test_exp, tar_exp):
             if flag == True:
                 return True, False, test_exp, tar_exp
             else:
-                return False, False, test_exp, tar_exp
+                flag = True
+                test_solves_list = list(test_solves.values())
+                target_solvers_list = list(tar_solves.values())
+                t1 = sorted(test_solves_list)
+                t2 = sorted(target_solvers_list)
+                for v1, v2 in zip(t1, t2):
+                    if abs(v1 - v2) > 1e-4:
+                        flag = False
+                        break
+                if flag:
+                    return True, False, test_exp, tar_exp
+                else:
+                    return False, False, test_exp, tar_exp
         except:
             return False, False, test_exp, tar_exp
         return False, False, test_exp, tar_exp
@@ -439,9 +463,9 @@ def result_multi(self, test_exp, tar_exp):
     def _compute_prefix_expression(self, pre_fix):
         st = list()
         operators = ["+", "-", "^", "*", "/"]
-        pre_fix = copy.deepcopy(pre_fix)
-        pre_fix.reverse()
-        for p in pre_fix:
+        pre_fix_ = copy.deepcopy(pre_fix)
+        pre_fix_.reverse()
+        for p in pre_fix_:
             if p not in operators:
                 pos = re.search("\d+\(", p)
                 if pos:
@@ -488,9 +512,9 @@ def _compute_prefix_expression_multi(self, pre_fix):
         st = list()
         operators = ["+", "-", "^", "*", "/", "=", "<BRG>"]
         unk_symbols = {}
-        pre_fix = copy.deepcopy(pre_fix)
-        pre_fix.reverse()
-        for p in pre_fix:
+        pre_fix_ = copy.deepcopy(pre_fix)
+        pre_fix_.reverse()
+        for p in pre_fix_:
             if p not in operators:
                 pos = re.search("\d+\(", p)
                 if pos:
@@ -662,7 +686,19 @@ def result_multi(self, test_exp, tar_exp):
             if flag == True:
                 return True, False, test_exp, tar_exp
             else:
-                return False, False, test_exp, tar_exp
+                flag = True
+                test_solves_list = list(test_solves.values())
+                target_solvers_list = list(tar_solves.values())
+                t1 = sorted(test_solves_list)
+                t2 = sorted(target_solvers_list)
+                for v1, v2 in zip(t1, t2):
+                    if abs(v1 - v2) > 1e-4:
+                        flag = False
+                        break
+                if flag:
+                    return True, False, test_exp, tar_exp
+                else:
+                    return False, False, test_exp, tar_exp
         except:
             return False, False, test_exp, tar_exp
         return False, False, test_exp, tar_exp
@@ -883,7 +919,19 @@ def result_multi(self, test_exp, tar_exp):
             if flag == True:
                 return True, False, tar_exp, tar_exp
             else:
-                return False, False, tar_exp, tar_exp
+                flag = True
+                test_solves_list = list(test_solves.values())
+                target_solvers_list = list(tar_solves.values())
+                t1 = sorted(test_solves_list)
+                t2 = sorted(target_solvers_list)
+                for v1, v2 in zip(t1, t2):
+                    if abs(v1 - v2) > 1e-4:
+                        flag = False
+                        break
+                if flag:
+                    return True, False, test_exp, tar_exp
+                else:
+                    return False, False, test_exp, tar_exp
         except:
             return False, False, tar_exp, tar_exp
 
@@ -1118,7 +1166,19 @@ def prefix_result_multi(self, test_exp, tar_exp):
             if flag == True:
                 return True, False, test_exp, tar_exp
             else:
-                return False, False, test_exp, tar_exp
+                flag = True
+                test_solves_list = list(test_solves.values())
+                target_solvers_list = list(tar_solves.values())
+                t1 = sorted(test_solves_list)
+                t2 = sorted(target_solvers_list)
+                for v1, v2 in zip(t1, t2):
+                    if abs(v1 - v2) > 1e-4:
+                        flag = False
+                        break
+                if flag:
+                    return True, False, test_exp, tar_exp
+                else:
+                    return False, False, test_exp, tar_exp
         except:
             return False, False, test_exp, tar_exp
         return False, False, test_exp, tar_exp
@@ -1224,7 +1284,19 @@ def postfix_result_multi(self, test_exp, tar_exp):
             if flag == True:
                 return True, False, test_exp, tar_exp
             else:
-                return False, False, test_exp, tar_exp
+                flag = True
+                test_solves_list = list(test_solves.values())
+                target_solvers_list = list(tar_solves.values())
+                t1 = sorted(test_solves_list)
+                t2 = sorted(target_solvers_list)
+                for v1, v2 in zip(t1, t2):
+                    if abs(v1 - v2) > 1e-4:
+                        flag = False
+                        break
+                if flag:
+                    return True, False, test_exp, tar_exp
+                else:
+                    return False, False, test_exp, tar_exp
         except:
             return False, False, test_exp, tar_exp
         return False, False, test_exp, tar_exp
diff --git a/mwptoolkit/loss/__init__.py b/mwptoolkit/loss/__init__.py
index d90f578..741c296 100644
--- a/mwptoolkit/loss/__init__.py
+++ b/mwptoolkit/loss/__init__.py
@@ -1,4 +1,12 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.loss import abstract_loss,binary_cross_entropy_loss,cross_entropy_loss,masked_cross_entropy_loss,mse_loss,nll_loss,smoothed_cross_entropy_loss
\ No newline at end of file
+from mwptoolkit.loss import abstract_loss,binary_cross_entropy_loss,cross_entropy_loss,masked_cross_entropy_loss,mse_loss,nll_loss,smoothed_cross_entropy_loss
+
+from mwptoolkit.loss.abstract_loss import AbstractLoss
+from mwptoolkit.loss.binary_cross_entropy_loss import BinaryCrossEntropyLoss
+from mwptoolkit.loss.cross_entropy_loss import CrossEntropyLoss
+from mwptoolkit.loss.masked_cross_entropy_loss import MaskedCrossEntropyLoss
+from mwptoolkit.loss.mse_loss import MSELoss
+from mwptoolkit.loss.nll_loss import NLLLoss
+from mwptoolkit.loss.smoothed_cross_entropy_loss import SmoothCrossEntropyLoss
diff --git a/mwptoolkit/model/Graph2Tree/__init__.py b/mwptoolkit/model/Graph2Tree/__init__.py
index 15ad28b..82eedbf 100644
--- a/mwptoolkit/model/Graph2Tree/__init__.py
+++ b/mwptoolkit/model/Graph2Tree/__init__.py
@@ -1,4 +1,7 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.model.Graph2Tree import graph2tree,multiencdec
\ No newline at end of file
+from mwptoolkit.model.Graph2Tree import graph2tree,multiencdec
+
+from mwptoolkit.model.Graph2Tree.graph2tree import Graph2Tree
+from mwptoolkit.model.Graph2Tree.multiencdec import MultiEncDec
diff --git a/mwptoolkit/model/PreTrain/__init__.py b/mwptoolkit/model/PreTrain/__init__.py
index bb27524..7d6e2c5 100644
--- a/mwptoolkit/model/PreTrain/__init__.py
+++ b/mwptoolkit/model/PreTrain/__init__.py
@@ -2,4 +2,8 @@
 from __future__ import print_function
 from __future__ import division
 
-from mwptoolkit.model.PreTrain import bertgen,gpt2,robertagen
\ No newline at end of file
+from mwptoolkit.model.PreTrain import bertgen,gpt2,robertagen
+
+from mwptoolkit.model.PreTrain.bertgen import BERTGen
+from mwptoolkit.model.PreTrain.gpt2 import GPT2
+from mwptoolkit.model.PreTrain.robertagen import RobertaGen
\ No newline at end of file
diff --git a/mwptoolkit/model/Seq2Seq/__init__.py b/mwptoolkit/model/Seq2Seq/__init__.py
index 8f07013..1edf3a9 100644
--- a/mwptoolkit/model/Seq2Seq/__init__.py
+++ b/mwptoolkit/model/Seq2Seq/__init__.py
@@ -1,4 +1,14 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.model.Seq2Seq import dns,ept,groupatt,lstm,mathen,rnnencdec,saligned,transformer,rnnvae
\ No newline at end of file
+from mwptoolkit.model.Seq2Seq import dns,ept,groupatt,lstm,mathen,rnnencdec,saligned,transformer,rnnvae
+
+from mwptoolkit.model.Seq2Seq.dns import DNS
+from mwptoolkit.model.Seq2Seq.ept import EPT
+from mwptoolkit.model.Seq2Seq.groupatt import GroupATT
+from mwptoolkit.model.Seq2Seq.lstm import LSTM
+from mwptoolkit.model.Seq2Seq.mathen import MathEN
+from mwptoolkit.model.Seq2Seq.rnnencdec import RNNEncDec
+from mwptoolkit.model.Seq2Seq.rnnvae import RNNVAE
+from mwptoolkit.model.Seq2Seq.saligned import Saligned
+from mwptoolkit.model.Seq2Seq.transformer import Transformer
\ No newline at end of file
diff --git a/mwptoolkit/model/Seq2Seq/mathen.py b/mwptoolkit/model/Seq2Seq/mathen.py
index adfd9cb..9fe1825 100644
--- a/mwptoolkit/model/Seq2Seq/mathen.py
+++ b/mwptoolkit/model/Seq2Seq/mathen.py
@@ -25,7 +25,7 @@ class MathEN(nn.Module):
         Wang et al. "Translating a Math Word Problem to a Expression Tree" in EMNLP 2018.
     """
     def __init__(self, config, dataset):
-        super(MatnEN, self).__init__()
+        super(MathEN, self).__init__()
         self.bidirectional = config["bidirectional"]
         self.embedding_size = config["embedding_size"]
         self.hidden_size = config["hidden_size"]
diff --git a/mwptoolkit/model/Seq2Tree/__init__.py b/mwptoolkit/model/Seq2Tree/__init__.py
index 43a9111..4536737 100644
--- a/mwptoolkit/model/Seq2Tree/__init__.py
+++ b/mwptoolkit/model/Seq2Tree/__init__.py
@@ -1,4 +1,10 @@
 from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
-from mwptoolkit.model.Seq2Tree import gts,sausolver,treelstm,trnn,tsn
\ No newline at end of file
+from mwptoolkit.model.Seq2Tree import gts,sausolver,treelstm,trnn,tsn
+
+from mwptoolkit.model.Seq2Tree.gts import GTS
+from mwptoolkit.model.Seq2Tree.sausolver import SAUSolver
+from mwptoolkit.model.Seq2Tree.treelstm import TreeLSTM
+from mwptoolkit.model.Seq2Tree.trnn import TRNN
+from mwptoolkit.model.Seq2Tree.tsn import TSN
diff --git a/mwptoolkit/model/Seq2Tree/sausolver.py b/mwptoolkit/model/Seq2Tree/sausolver.py
index 2681ad1..21f514d 100644
--- a/mwptoolkit/model/Seq2Tree/sausolver.py
+++ b/mwptoolkit/model/Seq2Tree/sausolver.py
@@ -12,7 +12,7 @@
 from mwptoolkit.module.Embedder.basic_embedder import BaiscEmbedder
 from mwptoolkit.module.Decoder.tree_decoder import SARTreeDecoder
 from mwptoolkit.module.Layer.tree_layers import NodeGenerater, SubTreeMerger, TreeNode, TreeEmbedding
-from mwptoolkit.module.Layer.tree_layers import Prediction, GenerateNode, Merge
+from mwptoolkit.module.Layer.tree_layers import Prediction, GenerateNode, Merge, SemanticAlignmentModule
 from mwptoolkit.module.Strategy.beam_search import TreeBeam
 from mwptoolkit.loss.masked_cross_entropy_loss import MaskedCrossEntropyLoss, masked_cross_entropy
 from mwptoolkit.loss.mse_loss import MSELoss
@@ -25,9 +25,10 @@ class SAUSolver(nn.Module):
     Reference:
         Qin et al. "Semantically-Aligned Universal Tree-Structured Solver for Math Word Problems" in EMNLP 2020.
     """
+
     def __init__(self, config, dataset):
-        super(SAUSolver,self).__init__()
-        #parameter
+        super(SAUSolver, self).__init__()
+        # parameter
         self.hidden_size = config["hidden_size"]
         self.device = config["device"]
         self.USE_CUDA = True if self.device == torch.device('cuda') else False
@@ -62,22 +63,26 @@ def __init__(self, config, dataset):
             self.out_pad_token = self.out_symbol2idx[SpecialTokens.PAD_TOKEN]
         except:
             self.out_pad_token = None
-        #module
+        # module
         self.embedder = BaiscEmbedder(self.vocab_size, self.embedding_size, self.dropout_ratio)
-        #self.t_encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type, self.dropout_ratio)
-        self.encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type, self.dropout_ratio, batch_first=False)
-        self.decoder = SARTreeDecoder(self.hidden_size, self.operator_nums, self.generate_size, self.dropout_ratio)
-        self.node_generater = GenerateNode(self.hidden_size, self.operator_nums, self.embedding_size, self.dropout_ratio)
+        # self.t_encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type, self.dropout_ratio)
+        self.encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type,
+                                       self.dropout_ratio, batch_first=False)
+        #self.decoder = SARTreeDecoder(self.hidden_size, self.operator_nums, self.generate_size, self.dropout_ratio)
+        self.decoder = Prediction(self.hidden_size,self.operator_nums,self.generate_size,self.dropout_ratio)
+        self.node_generater = GenerateNode(self.hidden_size, self.operator_nums, self.embedding_size,
+                                           self.dropout_ratio)
         self.merge = Merge(self.hidden_size, self.embedding_size, self.dropout_ratio)
-
+        self.sa = SemanticAlignmentModule(self.hidden_size,self.hidden_size,self.hidden_size)
         self.loss1 = MaskedCrossEntropyLoss()
         #
+
     def calculate_loss(self, batch_data):
         """Finish forward-propagating, calculating loss and back-propagation.
-        
+
         Args:
             batch_data (dict): one batch data.
-        
+
         Returns:
             float: loss value.
         """
@@ -95,16 +100,16 @@ def calculate_loss(self, batch_data):
         # sequence mask for attention
         unk = self.unk_token
 
-        loss = self.train_tree(seq,seq_length,target,target_length,\
-            nums_stack,num_size,generate_nums,num_pos,unk,num_start)
+        loss = self.train_tree(seq, seq_length, target, target_length, \
+                               nums_stack, num_size, generate_nums, num_pos, unk, num_start)
         return loss
 
     def model_test(self, batch_data):
         """Model test.
-        
+
         Args:
             batch_data (dict): one batch data.
-        
+
         Returns:
             tuple(list,list): predicted equation, target equation.
         """
@@ -120,33 +125,36 @@ def model_test(self, batch_data):
         generate_nums = self.generate_nums
         num_start = self.num_start
         # sequence mask for attention
-        all_node_output = self.evaluate_tree(seq, seq_length, generate_nums, num_pos, num_start, self.beam_size, self.max_out_len)
+        all_node_output = self.evaluate_tree(seq, seq_length, generate_nums, num_pos, num_start, self.beam_size,
+                                             self.max_out_len)
 
         all_output = self.convert_idx2symbol(all_node_output, num_list[0], copy_list(nums_stack[0]))
         targets = self.convert_idx2symbol(target[0], num_list[0], copy_list(nums_stack[0]))
         return all_output, targets
 
-    def train_tree(self, input_batch, input_length, target_batch, target_length, nums_stack_batch, num_size_batch, generate_nums, num_pos, unk, num_start, english=False):
+    def train_tree(self,input_batch, input_length, target_batch, target_length, nums_stack_batch, num_size_batch, generate_nums, num_pos, unk, num_start, 
+                                                   english=False,var_nums=[], batch_first=False):
         # sequence mask for attention
         seq_mask = []
         max_len = max(input_length)
         for i in input_length:
             seq_mask.append([0 for _ in range(i)] + [1 for _ in range(i, max_len)])
-        seq_mask = torch.BoolTensor(seq_mask)
+        seq_mask = torch.ByteTensor(seq_mask)
 
         num_mask = []
-        max_num_size = max(num_size_batch) + len(generate_nums)
+        max_num_size = max(num_size_batch) + len(generate_nums) + len(var_nums)  # 最大的位置列表数目+常识数字数目+未知数列表
         for i in num_size_batch:
-            d = i + len(generate_nums)
+            d = i + len(generate_nums) + len(var_nums)
             num_mask.append([0] * d + [1] * (max_num_size - d))
-        num_mask = torch.BoolTensor(num_mask)
+        num_mask = torch.ByteTensor(num_mask)  # 用于屏蔽无关数字，防止生成错误的Nx
+
+        #unk = output_lang.word2index["UNK"]
 
         # Turn padded arrays into (batch_size x max_len) tensors, transpose into (max_len x batch_size)
         input_var = input_batch.transpose(0, 1)
-
         target = target_batch.transpose(0, 1)
 
-        padding_hidden = torch.FloatTensor([0.0 for _ in range(self.hidden_size)]).unsqueeze(0)
+        padding_hidden = torch.FloatTensor([0.0 for _ in range(self.decoder.hidden_size)]).unsqueeze(0)
         batch_size = len(input_length)
 
         if self.USE_CUDA:
@@ -155,105 +163,120 @@ def train_tree(self, input_batch, input_length, target_batch, target_length, num
             padding_hidden = padding_hidden.cuda()
             num_mask = num_mask.cuda()
 
+        # Zero gradients of both optimizers
         # Run words through encoder
+
+        #encoder_outputs, problem_output = self.encoder(input_var, input_length)
         seq_emb = self.embedder(input_var)
         pade_outputs, _ = self.encoder(seq_emb, input_length)
         problem_output = pade_outputs[-1, :, :self.hidden_size] + pade_outputs[0, :, self.hidden_size:]
         encoder_outputs = pade_outputs[:, :, :self.hidden_size] + pade_outputs[:, :, self.hidden_size:]
         # Prepare input and output variables
-        node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]
+        node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]  # root embedding B x 1
 
         max_target_length = max(target_length)
 
         all_node_outputs = []
+        all_sa_outputs = []
         # all_leafs = []
-        sub_tree_outputs = []
-        sub_tree_target = []
-        sub_tree_mask = []
 
         copy_num_len = [len(_) for _ in num_pos]
         num_size = max(copy_num_len)
-        all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size, self.hidden_size)
+        # 提取与问题相关的数字embedding
+        all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size,
+                                                                  self.encoder.hidden_size)
+
+        embeddings_stacks = [[] for _ in range(batch_size)]  # B x 1  当前的tree state/ subtree embedding / output
+        left_childs = [None for _ in range(batch_size)]  # B x 1
 
-        embeddings_stacks = [[] for _ in range(batch_size)]
-        left_childs = [None for _ in range(batch_size)]
         for t in range(max_target_length):
-            num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(node_stacks, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden, seq_mask,
-                                                                                                       num_mask)
+            num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(
+                node_stacks, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden, seq_mask, num_mask)
 
             # all_leafs.append(p_leaf)
             outputs = torch.cat((op, num_score), 1)
             all_node_outputs.append(outputs)
 
-            target_t, generate_input = self.generate_tree_input(target[t].tolist(), outputs, nums_stack_batch, num_start, unk)
+            target_t, generate_input = self.generate_tree_input(target[t].tolist(), outputs, nums_stack_batch, num_start,
+                                                           unk)
             target[t] = target_t
             if self.USE_CUDA:
                 generate_input = generate_input.cuda()
             left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input, current_context)
             left_childs = []
-            sub_tree_emb = []
-            loss_mask = []
-            for idx, l, r, node_stack, i, o in zip(range(batch_size), left_child.split(1), right_child.split(1), node_stacks, target[t].tolist(), embeddings_stacks):
+            for idx, l, r, node_stack, i, o in zip(range(batch_size), left_child.split(1), right_child.split(1),
+                                                   node_stacks, target[t].tolist(), embeddings_stacks):
                 if len(node_stack) != 0:
                     node = node_stack.pop()
                 else:
                     left_childs.append(None)
-                    sub_tree_emb.append(padding_hidden)
-                    loss_mask.append(torch.zeros(1, dtype=torch.float))
                     continue
 
-                if i < num_start:
+                # 未知数当数字处理，SEP当操作符处理
+                if i < num_start:  # 非数字
                     node_stack.append(TreeNode(r))
                     node_stack.append(TreeNode(l, left_flag=True))
-                    o.append(TreeEmbedding(node_label[idx].unsqueeze(0), False))
-                else:
+                    o.append(TreeEmbedding(node_label[idx].unsqueeze(0), terminal=False))
+                    # print(o[-1].embedding.size())
+                    # print(encoder_outputs[idx].size())
+                else:  # 数字
                     current_num = current_nums_embeddings[idx, i - num_start].unsqueeze(0)
                     while len(o) > 0 and o[-1].terminal:
                         sub_stree = o.pop()
                         op = o.pop()
-                        current_num = self.merge(op.embedding, sub_stree.embedding, current_num)
-                    o.append(TreeEmbedding(current_num, True))
+                        current_num = self.merge(op.embedding, sub_stree.embedding, current_num)  # Subtree embedding
+                        if batch_first:
+                            encoder_mapping, decoder_mapping = self.sa(current_num, encoder_outputs[idx])
+                        else:
+                            temp_encoder_outputs = encoder_outputs.transpose(0, 1)
+                            encoder_mapping, decoder_mapping = self.sa(current_num,temp_encoder_outputs[idx])
+                        all_sa_outputs.append((encoder_mapping, decoder_mapping))
+                    o.append(TreeEmbedding(current_num, terminal=True))
+
                 if len(o) > 0 and o[-1].terminal:
                     left_childs.append(o[-1].embedding)
-                    sub_tree_emb.append(o[-1].embedding)
-                    loss_mask.append(torch.ones(1, dtype=torch.float))
+
                 else:
                     left_childs.append(None)
-                    sub_tree_emb.append(padding_hidden)
-                    loss_mask.append(torch.zeros(1, dtype=torch.float))
-
-            sub_tree_emb = torch.stack(sub_tree_emb)
-            loss_mask = torch.stack(loss_mask)
-            score = self.decoder.attn(sub_tree_emb.transpose(0, 1), encoder_outputs, seq_mask)
-            s_aligned_vector = score.bmm(encoder_outputs.transpose(0, 1))
-            s_aligned_a, s_aligned_d = self.decoder.Semantically_Aligned_Regularization(sub_tree_emb, s_aligned_vector)
-            sub_tree_outputs.append(s_aligned_a)
-            sub_tree_target.append(s_aligned_d)
-            sub_tree_mask.append(loss_mask)
 
         # all_leafs = torch.stack(all_leafs, dim=1)  # B x S x 2
         all_node_outputs = torch.stack(all_node_outputs, dim=1)  # B x S x N
-        sub_tree_outputs = torch.cat(sub_tree_outputs, dim=1)
-        sub_tree_target = torch.cat(sub_tree_target, dim=1)
-        sub_tree_mask = torch.cat(sub_tree_mask, dim=1)
 
-        target = target.transpose(0, 1).contiguous()
+        target = target.transpose(0, 1).contiguous()  # B x S
+
         if self.USE_CUDA:
             # all_leafs = all_leafs.cuda()
             all_node_outputs = all_node_outputs.cuda()
             target = target.cuda()
+            new_all_sa_outputs = []
+            for sa_pair in all_sa_outputs:
+                new_all_sa_outputs.append((sa_pair[0].cuda(), sa_pair[1].cuda()))
+            all_sa_outputs = new_all_sa_outputs
             target_length = torch.LongTensor(target_length).cuda()
         else:
             target_length = torch.LongTensor(target_length)
 
+        semantic_alignment_loss = nn.MSELoss()
+        total_semanti_alognment_loss = 0
+        sa_len = len(all_sa_outputs)
+        for sa_pair in all_sa_outputs:
+            total_semanti_alognment_loss += semantic_alignment_loss(sa_pair[0], sa_pair[1])
+        # print(total_semanti_alognment_loss)
+        total_semanti_alognment_loss = total_semanti_alognment_loss / sa_len
+        # print(total_semanti_alognment_loss)
+
         # op_target = target < num_start
         # loss_0 = masked_cross_entropy_without_logit(all_leafs, op_target.long(), target_length)
-        loss_1 = masked_cross_entropy(all_node_outputs, target, target_length)
-        loss_2 = self.mse_loss(sub_tree_outputs, sub_tree_target, sub_tree_mask)
-        #self.loss2.eval_batch(sub_tree_outputs, sub_tree_target)
-        loss = loss_1 + self.loss_weight * loss_2
+        loss = masked_cross_entropy(all_node_outputs, target,target_length) + 0.01 * total_semanti_alognment_loss
+        # loss = loss_0 + loss_1
         loss.backward()
-        return loss.item()
+        # clip the grad
+        # torch.nn.utils.clip_grad_norm_(encoder.parameters(), 5)
+        # torch.nn.utils.clip_grad_norm_(predict.parameters(), 5)
+        # torch.nn.utils.clip_grad_norm_(generate.parameters(), 5)
+
+        # Update parameters with optimizers
+        return loss.item()  # , loss_0.item(), loss_1.item()
 
     def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_start, beam_size=5, max_length=30):
 
@@ -283,7 +306,8 @@ def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_s
         node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]
 
         num_size = len(num_pos[0])
-        all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size, self.hidden_size)
+        all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size,
+                                                                       self.hidden_size)
         # B x P x N
         embeddings_stacks = [[] for _ in range(batch_size)]
         left_childs = [None for _ in range(batch_size)]
@@ -300,8 +324,13 @@ def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_s
                 # left_childs = torch.stack(b.left_childs)
                 left_childs = b.left_childs
 
-                num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(b.node_stack, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden,
-                                                                                                           seq_mask, num_mask)
+                num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(b.node_stack,
+                                                                                                           left_childs,
+                                                                                                           encoder_outputs,
+                                                                                                           all_nums_encoder_outputs,
+                                                                                                           padding_hidden,
+                                                                                                           seq_mask,
+                                                                                                           num_mask)
 
                 out_score = nn.functional.log_softmax(torch.cat((op, num_score), dim=1), dim=1)
 
@@ -324,7 +353,8 @@ def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_s
                         generate_input = torch.LongTensor([out_token])
                         if self.USE_CUDA:
                             generate_input = generate_input.cuda()
-                        left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input, current_context)
+                        left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input,
+                                                                                  current_context)
 
                         current_node_stack[0].append(TreeNode(right_child))
                         current_node_stack[0].append(TreeNode(left_child, left_flag=True))
@@ -342,7 +372,8 @@ def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_s
                         current_left_childs.append(current_embeddings_stacks[0][-1].embedding)
                     else:
                         current_left_childs.append(None)
-                    current_beams.append(TreeBeam(b.score + float(tv), current_node_stack, current_embeddings_stacks, current_left_childs, current_out))
+                    current_beams.append(TreeBeam(b.score + float(tv), current_node_stack, current_embeddings_stacks,
+                                                  current_left_childs, current_out))
             beams = sorted(current_beams, key=lambda x: x.score, reverse=True)
             beams = beams[:beam_size]
             flag = True
@@ -354,6 +385,175 @@ def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_s
 
         return beams[0].out
 
+    # def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_start, beam_size=5, max_length=30,var_nums=[]):
+    #     # sequence mask for attention
+    #     seq_mask = torch.ByteTensor(1, input_length).fill_(0)
+    #     # Turn padded arrays into (batch_size x max_len) tensors, transpose into (max_len x batch_size)
+    #     input_var = torch.LongTensor(input_batch).unsqueeze(1)
+    # 
+    #     num_mask = torch.ByteTensor(1, len(num_pos) + len(generate_nums) + len(var_nums)).fill_(0)
+    # 
+    #     # Set to not-training mode to disable dropout
+    # 
+    #     padding_hidden = torch.FloatTensor([0.0 for _ in range(self.decoder.hidden_size)]).unsqueeze(0)
+    # 
+    #     batch_size = 1
+    # 
+    #     if self.USE_CUDA:
+    #         input_var = input_var.cuda()
+    #         seq_mask = seq_mask.cuda()
+    #         padding_hidden = padding_hidden.cuda()
+    #         num_mask = num_mask.cuda()
+    # 
+    #     # Run words through encoder
+    #     encoder_outputs, problem_output = self.encoder(input_var, input_length)
+    # 
+    #     # Prepare input and output variables  # # root embedding B x 1
+    #     node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]
+    # 
+    #     num_size = len(num_pos)
+    #     # 提取与问题相关的数字embedding
+    #     all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, [num_pos], batch_size, num_size,
+    #                                                               self.encoder.hidden_size)
+    #     # B x P x N
+    #     embeddings_stacks = [[] for _ in range(batch_size)]
+    #     left_childs = [None for _ in range(batch_size)]
+    #     beam_search=True
+    #     if beam_search:
+    #         beams = [TreeBeam(0.0, node_stacks, embeddings_stacks, left_childs, [])]
+    # 
+    #         for t in range(max_length):
+    #             current_beams = []
+    #             while len(beams) > 0:
+    #                 b = beams.pop()
+    #                 if len(b.node_stack[0]) == 0:
+    #                     current_beams.append(b)
+    #                     continue
+    #                 # left_childs = torch.stack(b.left_childs)
+    #                 left_childs = b.left_childs
+    # 
+    #                 num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(
+    #                     b.node_stack, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden,
+    #                     seq_mask, num_mask)
+    # 
+    #                 # leaf = p_leaf[:, 0].unsqueeze(1)
+    #                 # repeat_dims = [1] * leaf.dim()
+    #                 # repeat_dims[1] = op.size(1)
+    #                 # leaf = leaf.repeat(*repeat_dims)
+    #                 #
+    #                 # non_leaf = p_leaf[:, 1].unsqueeze(1)
+    #                 # repeat_dims = [1] * non_leaf.dim()
+    #                 # repeat_dims[1] = num_score.size(1)
+    #                 # non_leaf = non_leaf.repeat(*repeat_dims)
+    #                 #
+    #                 # p_leaf = torch.cat((leaf, non_leaf), dim=1)
+    #                 out_score = nn.functional.log_softmax(torch.cat((op, num_score), dim=1), dim=1)
+    # 
+    #                 # out_score = p_leaf * out_score
+    # 
+    #                 topv, topi = out_score.topk(beam_size)
+    # 
+    #                 # is_leaf = int(topi[0])
+    #                 # if is_leaf:
+    #                 #     topv, topi = op.topk(1)
+    #                 #     out_token = int(topi[0])
+    #                 # else:
+    #                 #     topv, topi = num_score.topk(1)
+    #                 #     out_token = int(topi[0]) + num_start
+    #                 for tv, ti in zip(topv.split(1, dim=1), topi.split(1, dim=1)):
+    #                     current_node_stack = copy_list(b.node_stack)
+    #                     current_left_childs = []
+    #                     current_embeddings_stacks = copy_list(b.embedding_stack)
+    #                     current_out = copy.deepcopy(b.out)
+    #                     out_token = int(ti)
+    #                     current_out.append(out_token)
+    # 
+    #                     node = current_node_stack[0].pop()
+    # 
+    #                     # var_num当时数字处理，SEP/;当操作符处理
+    #                     if out_token < num_start:  # 非数字
+    #                         generate_input = torch.LongTensor([out_token])
+    #                         if self.USE_CUDA:
+    #                             generate_input = generate_input.cuda()
+    #                         left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input,
+    #                                                                        current_context)
+    # 
+    #                         current_node_stack[0].append(TreeNode(right_child))
+    #                         current_node_stack[0].append(TreeNode(left_child, left_flag=True))
+    # 
+    #                         current_embeddings_stacks[0].append(TreeEmbedding(node_label[0].unsqueeze(0), False))
+    #                     else:  # 数字
+    #                         current_num = current_nums_embeddings[0, out_token - num_start].unsqueeze(0)
+    # 
+    #                         while len(current_embeddings_stacks[0]) > 0 and current_embeddings_stacks[0][-1].terminal:
+    #                             sub_stree = current_embeddings_stacks[0].pop()
+    #                             op = current_embeddings_stacks[0].pop()
+    #                             current_num = self.merge(op.embedding, sub_stree.embedding, current_num)
+    #                         current_embeddings_stacks[0].append(TreeEmbedding(current_num, True))
+    #                     if len(current_embeddings_stacks[0]) > 0 and current_embeddings_stacks[0][-1].terminal:
+    #                         current_left_childs.append(current_embeddings_stacks[0][-1].embedding)
+    #                     else:
+    #                         current_left_childs.append(None)
+    #                     current_beams.append(
+    #                         TreeBeam(b.score + float(tv), current_node_stack, current_embeddings_stacks,
+    #                                  current_left_childs, current_out))
+    #             beams = sorted(current_beams, key=lambda x: x.score, reverse=True)
+    #             beams = beams[:beam_size]
+    #             flag = True
+    #             for b in beams:
+    #                 if len(b.node_stack[0]) != 0:
+    #                     flag = False
+    #             if flag:
+    #                 break
+    # 
+    #         return beams[0].out
+    #     else:
+    #         all_node_outputs = []
+    #         for t in range(max_length):
+    #             num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(
+    #                 node_stacks, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden,
+    #                 seq_mask, num_mask)
+    # 
+    #             out_scores = nn.functional.log_softmax(torch.cat((op, num_score), dim=1), dim=1)
+    #             out_tokens = torch.argmax(out_scores, dim=1)  # B
+    #             all_node_outputs.append(out_tokens)
+    #             left_childs = []
+    #             for idx, node_stack, out_token, embeddings_stack in zip(range(batch_size), node_stacks, out_tokens,
+    #                                                                     embeddings_stacks):
+    #                 # node = node_stack.pop()
+    #                 if len(node_stack) != 0:
+    #                     node = node_stack.pop()
+    #                 else:
+    #                     left_childs.append(None)
+    #                     continue
+    #                 # var_num当时数字处理，SEP/;当操作符处理
+    #                 if out_token < num_start:  # 非数字
+    #                     generate_input = torch.LongTensor([out_token])
+    #                     if self.USE_CUDA:
+    #                         generate_input = generate_input.cuda()
+    #                     left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input,
+    #                                                                    current_context)
+    #                     node_stack.append(TreeNode(right_child))
+    #                     node_stack.append(TreeNode(left_child, left_flag=True))
+    #                     embeddings_stack.append(TreeEmbedding(node_label.unsqueeze(0), False))
+    #                 else:  # 数字
+    #                     current_num = current_nums_embeddings[idx, out_token - num_start].unsqueeze(0)
+    #                     while len(embeddings_stack) > 0 and embeddings_stack[-1].terminal:
+    #                         sub_stree = embeddings_stack.pop()
+    #                         op = embeddings_stack.pop()
+    #                         current_num = self.merge(op.embedding.squeeze(0), sub_stree.embedding, current_num)
+    #                     embeddings_stack.append(TreeEmbedding(current_num, terminal=True))
+    # 
+    #                 if len(embeddings_stack) > 0 and embeddings_stack[-1].terminal:
+    #                     left_childs.append(embeddings_stack[-1].embedding)
+    #                 else:
+    #                     left_childs.append(None)
+    # 
+    #         # all_leafs = torch.stack(all_leafs, dim=1)  # B x S x 2
+    #         all_node_outputs = torch.stack(all_node_outputs, dim=1)  # B x S x N
+    #         all_node_outputs = all_node_outputs.cpu().numpy()
+    #         return all_node_outputs[0]
+
     def get_all_number_encoder_outputs(self, encoder_outputs, num_pos, batch_size, num_size, hidden_size):
         indices = list()
         sen_len = encoder_outputs.size(0)
@@ -403,7 +603,7 @@ def mse_loss(self, outputs, targets, mask=None):
         return torch.sum(y)
 
     def convert_idx2symbol(self, output, num_list, num_stack):
-        #batch_size=output.size(0)
+        # batch_size=output.size(0)
         '''batch_size=1'''
         seq_len = len(output)
         num_len = len(num_list)
@@ -433,6 +633,420 @@ def convert_idx2symbol(self, output, num_list, num_stack):
         return output_list
 
 
+# class SAUSolver(nn.Module):
+#     """
+#     Reference:
+#         Qin et al. "Semantically-Aligned Universal Tree-Structured Solver for Math Word Problems" in EMNLP 2020.
+#     """
+#     def __init__(self, config, dataset):
+#         super(SAUSolver,self).__init__()
+#         #parameter
+#         self.hidden_size = config["hidden_size"]
+#         self.device = config["device"]
+#         self.USE_CUDA = True if self.device == torch.device('cuda') else False
+#         self.beam_size = config['beam_size']
+#         self.max_out_len = config['max_output_len']
+#         self.embedding_size = config["embedding_size"]
+#         self.dropout_ratio = config["dropout_ratio"]
+#         self.num_layers = config["num_layers"]
+#         self.rnn_cell_type = config["rnn_cell_type"]
+#         self.loss_weight = config['loss_weight']
+# 
+#         self.vocab_size = len(dataset.in_idx2word)
+#         self.out_symbol2idx = dataset.out_symbol2idx
+#         self.out_idx2symbol = dataset.out_idx2symbol
+#         generate_list = dataset.generate_list
+#         self.generate_nums = [self.out_symbol2idx[symbol] for symbol in generate_list]
+#         self.mask_list = NumMask.number
+#         self.num_start = dataset.num_start
+#         self.operator_nums = dataset.operator_nums
+#         self.generate_size = len(generate_list)
+# 
+#         self.unk_token = self.out_symbol2idx[SpecialTokens.UNK_TOKEN]
+#         try:
+#             self.out_sos_token = self.out_symbol2idx[SpecialTokens.SOS_TOKEN]
+#         except:
+#             self.out_sos_token = None
+#         try:
+#             self.out_eos_token = self.out_symbol2idx[SpecialTokens.EOS_TOKEN]
+#         except:
+#             self.out_eos_token = None
+#         try:
+#             self.out_pad_token = self.out_symbol2idx[SpecialTokens.PAD_TOKEN]
+#         except:
+#             self.out_pad_token = None
+#         #module
+#         self.embedder = BaiscEmbedder(self.vocab_size, self.embedding_size, self.dropout_ratio)
+#         #self.t_encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type, self.dropout_ratio)
+#         self.encoder = BasicRNNEncoder(self.embedding_size, self.hidden_size, self.num_layers, self.rnn_cell_type, self.dropout_ratio, batch_first=False)
+#         self.decoder = SARTreeDecoder(self.hidden_size, self.operator_nums, self.generate_size, self.dropout_ratio)
+#         self.node_generater = GenerateNode(self.hidden_size, self.operator_nums, self.embedding_size, self.dropout_ratio)
+#         self.merge = Merge(self.hidden_size, self.embedding_size, self.dropout_ratio)
+# 
+#         self.loss1 = MaskedCrossEntropyLoss()
+#         #
+#     def calculate_loss(self, batch_data):
+#         """Finish forward-propagating, calculating loss and back-propagation.
+#         
+#         Args:
+#             batch_data (dict): one batch data.
+#         
+#         Returns:
+#             float: loss value.
+#         """
+#         seq = batch_data["question"]
+#         seq_length = batch_data["ques len"]
+#         nums_stack = batch_data["num stack"]
+#         num_size = batch_data["num size"]
+#         num_pos = batch_data["num pos"]
+#         target = batch_data["equation"]
+#         target_length = batch_data["equ len"]
+#         equ_mask = batch_data["equ mask"]
+#         num_list = batch_data['num list']
+#         generate_nums = self.generate_nums
+#         num_start = self.num_start
+#         # sequence mask for attention
+#         unk = self.unk_token
+# 
+#         loss = self.train_tree(seq,seq_length,target,target_length,\
+#             nums_stack,num_size,generate_nums,num_pos,unk,num_start)
+#         return loss
+# 
+#     def model_test(self, batch_data):
+#         """Model test.
+#         
+#         Args:
+#             batch_data (dict): one batch data.
+#         
+#         Returns:
+#             tuple(list,list): predicted equation, target equation.
+#         """
+#         seq = batch_data["question"]
+#         seq_length = batch_data["ques len"]
+#         nums_stack = batch_data["num stack"]
+#         num_size = batch_data["num size"]
+#         num_pos = batch_data["num pos"]
+#         target = batch_data["equation"]
+#         target_length = batch_data["equ len"]
+#         equ_mask = batch_data["equ mask"]
+#         num_list = batch_data['num list']
+#         generate_nums = self.generate_nums
+#         num_start = self.num_start
+#         # sequence mask for attention
+#         all_node_output = self.evaluate_tree(seq, seq_length, generate_nums, num_pos, num_start, self.beam_size, self.max_out_len)
+# 
+#         all_output = self.convert_idx2symbol(all_node_output, num_list[0], copy_list(nums_stack[0]))
+#         targets = self.convert_idx2symbol(target[0], num_list[0], copy_list(nums_stack[0]))
+#         return all_output, targets
+# 
+#     def train_tree(self, input_batch, input_length, target_batch, target_length, nums_stack_batch, num_size_batch, generate_nums, num_pos, unk, num_start, english=False):
+#         # sequence mask for attention
+#         seq_mask = []
+#         max_len = max(input_length)
+#         for i in input_length:
+#             seq_mask.append([0 for _ in range(i)] + [1 for _ in range(i, max_len)])
+#         seq_mask = torch.BoolTensor(seq_mask)
+# 
+#         num_mask = []
+#         max_num_size = max(num_size_batch) + len(generate_nums)
+#         for i in num_size_batch:
+#             d = i + len(generate_nums)
+#             num_mask.append([0] * d + [1] * (max_num_size - d))
+#         num_mask = torch.BoolTensor(num_mask)
+# 
+#         # Turn padded arrays into (batch_size x max_len) tensors, transpose into (max_len x batch_size)
+#         input_var = input_batch.transpose(0, 1)
+# 
+#         target = target_batch.transpose(0, 1)
+# 
+#         padding_hidden = torch.FloatTensor([0.0 for _ in range(self.hidden_size)]).unsqueeze(0)
+#         batch_size = len(input_length)
+# 
+#         if self.USE_CUDA:
+#             input_var = input_var.cuda()
+#             seq_mask = seq_mask.cuda()
+#             padding_hidden = padding_hidden.cuda()
+#             num_mask = num_mask.cuda()
+# 
+#         # Run words through encoder
+#         seq_emb = self.embedder(input_var)
+#         pade_outputs, _ = self.encoder(seq_emb, input_length)
+#         problem_output = pade_outputs[-1, :, :self.hidden_size] + pade_outputs[0, :, self.hidden_size:]
+#         encoder_outputs = pade_outputs[:, :, :self.hidden_size] + pade_outputs[:, :, self.hidden_size:]
+#         # Prepare input and output variables
+#         node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]
+# 
+#         max_target_length = max(target_length)
+# 
+#         all_node_outputs = []
+#         # all_leafs = []
+#         sub_tree_outputs = []
+#         sub_tree_target = []
+#         sub_tree_mask = []
+# 
+#         copy_num_len = [len(_) for _ in num_pos]
+#         num_size = max(copy_num_len)
+#         all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size, self.hidden_size)
+# 
+#         embeddings_stacks = [[] for _ in range(batch_size)]
+#         left_childs = [None for _ in range(batch_size)]
+#         for t in range(max_target_length):
+#             num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(node_stacks, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden, seq_mask,
+#                                                                                                        num_mask)
+# 
+#             # all_leafs.append(p_leaf)
+#             outputs = torch.cat((op, num_score), 1)
+#             all_node_outputs.append(outputs)
+# 
+#             target_t, generate_input = self.generate_tree_input(target[t].tolist(), outputs, nums_stack_batch, num_start, unk)
+#             target[t] = target_t
+#             if self.USE_CUDA:
+#                 generate_input = generate_input.cuda()
+#             left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input, current_context)
+#             left_childs = []
+#             sub_tree_emb = []
+#             loss_mask = []
+#             for idx, l, r, node_stack, i, o in zip(range(batch_size), left_child.split(1), right_child.split(1), node_stacks, target[t].tolist(), embeddings_stacks):
+#                 if len(node_stack) != 0:
+#                     node = node_stack.pop()
+#                 else:
+#                     left_childs.append(None)
+#                     sub_tree_emb.append(padding_hidden)
+#                     loss_mask.append(torch.zeros(1, dtype=torch.float))
+#                     continue
+# 
+#                 if i < num_start:
+#                     node_stack.append(TreeNode(r))
+#                     node_stack.append(TreeNode(l, left_flag=True))
+#                     o.append(TreeEmbedding(node_label[idx].unsqueeze(0), False))
+#                 else:
+#                     current_num = current_nums_embeddings[idx, i - num_start].unsqueeze(0)
+#                     while len(o) > 0 and o[-1].terminal:
+#                         sub_stree = o.pop()
+#                         op = o.pop()
+#                         current_num = self.merge(op.embedding, sub_stree.embedding, current_num)
+#                     o.append(TreeEmbedding(current_num, True))
+#                 if len(o) > 0 and o[-1].terminal:
+#                     left_childs.append(o[-1].embedding)
+#                     sub_tree_emb.append(o[-1].embedding)
+#                     loss_mask.append(torch.ones(1, dtype=torch.float))
+#                 else:
+#                     left_childs.append(None)
+#                     sub_tree_emb.append(padding_hidden)
+#                     loss_mask.append(torch.zeros(1, dtype=torch.float))
+# 
+#             sub_tree_emb = torch.stack(sub_tree_emb)
+#             loss_mask = torch.stack(loss_mask)
+#             #score = self.decoder.attn(sub_tree_emb.transpose(0, 1), encoder_outputs, seq_mask)
+#             score = self.decoder.saligned_attn(sub_tree_emb.transpose(0, 1), encoder_outputs, seq_mask)
+#             s_aligned_vector = score.bmm(encoder_outputs.transpose(0, 1)) #vector a in paper
+#             s_aligned_a, s_aligned_d = self.decoder.Semantically_Aligned_Regularization(sub_tree_emb, s_aligned_vector)
+#             sub_tree_outputs.append(s_aligned_a)
+#             sub_tree_target.append(s_aligned_d)
+#             sub_tree_mask.append(loss_mask)
+# 
+#         # all_leafs = torch.stack(all_leafs, dim=1)  # B x S x 2
+#         all_node_outputs = torch.stack(all_node_outputs, dim=1)  # B x S x N
+#         sub_tree_outputs = torch.cat(sub_tree_outputs, dim=1)
+#         sub_tree_target = torch.cat(sub_tree_target, dim=1)
+#         sub_tree_mask = torch.cat(sub_tree_mask, dim=1)
+# 
+#         target = target.transpose(0, 1).contiguous()
+#         if self.USE_CUDA:
+#             # all_leafs = all_leafs.cuda()
+#             all_node_outputs = all_node_outputs.cuda()
+#             target = target.cuda()
+#             target_length = torch.LongTensor(target_length).cuda()
+#         else:
+#             target_length = torch.LongTensor(target_length)
+# 
+#         # op_target = target < num_start
+#         # loss_0 = masked_cross_entropy_without_logit(all_leafs, op_target.long(), target_length)
+#         loss_1 = masked_cross_entropy(all_node_outputs, target, target_length)
+#         loss_2 = self.mse_loss(sub_tree_outputs, sub_tree_target, sub_tree_mask)
+#         #self.loss2.eval_batch(sub_tree_outputs, sub_tree_target)
+#         loss = loss_1 + self.loss_weight * loss_2
+#         loss.backward()
+#         return loss.item()
+# 
+#     def evaluate_tree(self, input_batch, input_length, generate_nums, num_pos, num_start, beam_size=5, max_length=30):
+# 
+#         seq_mask = torch.BoolTensor(1, input_length).fill_(0)
+#         # Turn padded arrays into (batch_size x max_len) tensors, transpose into (max_len x batch_size)
+#         input_var = input_batch.transpose(0, 1)
+# 
+#         num_mask = torch.BoolTensor(1, len(num_pos[0]) + len(generate_nums)).fill_(0)
+# 
+#         padding_hidden = torch.FloatTensor([0.0 for _ in range(self.hidden_size)]).unsqueeze(0)
+# 
+#         batch_size = 1
+# 
+#         if self.USE_CUDA:
+#             input_var = input_var.cuda()
+#             seq_mask = seq_mask.cuda()
+#             padding_hidden = padding_hidden.cuda()
+#             num_mask = num_mask.cuda()
+#         # Run words through encoder
+# 
+#         seq_emb = self.embedder(input_var)
+#         pade_outputs, _ = self.encoder(seq_emb, input_length)
+#         problem_output = pade_outputs[-1, :, :self.hidden_size] + pade_outputs[0, :, self.hidden_size:]
+#         encoder_outputs = pade_outputs[:, :, :self.hidden_size] + pade_outputs[:, :, self.hidden_size:]
+# 
+#         # Prepare input and output variables
+#         node_stacks = [[TreeNode(_)] for _ in problem_output.split(1, dim=0)]
+# 
+#         num_size = len(num_pos[0])
+#         all_nums_encoder_outputs = self.get_all_number_encoder_outputs(encoder_outputs, num_pos, batch_size, num_size, self.hidden_size)
+#         # B x P x N
+#         embeddings_stacks = [[] for _ in range(batch_size)]
+#         left_childs = [None for _ in range(batch_size)]
+# 
+#         beams = [TreeBeam(0.0, node_stacks, embeddings_stacks, left_childs, [])]
+# 
+#         for t in range(max_length):
+#             current_beams = []
+#             while len(beams) > 0:
+#                 b = beams.pop()
+#                 if len(b.node_stack[0]) == 0:
+#                     current_beams.append(b)
+#                     continue
+#                 # left_childs = torch.stack(b.left_childs)
+#                 left_childs = b.left_childs
+# 
+#                 num_score, op, current_embeddings, current_context, current_nums_embeddings = self.decoder(b.node_stack, left_childs, encoder_outputs, all_nums_encoder_outputs, padding_hidden,
+#                                                                                                            seq_mask, num_mask)
+# 
+#                 out_score = nn.functional.log_softmax(torch.cat((op, num_score), dim=1), dim=1)
+# 
+#                 # out_score = p_leaf * out_score
+# 
+#                 topv, topi = out_score.topk(beam_size)
+# 
+#                 for tv, ti in zip(topv.split(1, dim=1), topi.split(1, dim=1)):
+#                     current_node_stack = copy_list(b.node_stack)
+#                     current_left_childs = []
+#                     current_embeddings_stacks = copy_list(b.embedding_stack)
+#                     current_out = copy.deepcopy(b.out)
+# 
+#                     out_token = int(ti)
+#                     current_out.append(out_token)
+# 
+#                     node = current_node_stack[0].pop()
+# 
+#                     if out_token < num_start:
+#                         generate_input = torch.LongTensor([out_token])
+#                         if self.USE_CUDA:
+#                             generate_input = generate_input.cuda()
+#                         left_child, right_child, node_label = self.node_generater(current_embeddings, generate_input, current_context)
+# 
+#                         current_node_stack[0].append(TreeNode(right_child))
+#                         current_node_stack[0].append(TreeNode(left_child, left_flag=True))
+# 
+#                         current_embeddings_stacks[0].append(TreeEmbedding(node_label[0].unsqueeze(0), False))
+#                     else:
+#                         current_num = current_nums_embeddings[0, out_token - num_start].unsqueeze(0)
+# 
+#                         while len(current_embeddings_stacks[0]) > 0 and current_embeddings_stacks[0][-1].terminal:
+#                             sub_stree = current_embeddings_stacks[0].pop()
+#                             op = current_embeddings_stacks[0].pop()
+#                             current_num = self.merge(op.embedding, sub_stree.embedding, current_num)
+#                         current_embeddings_stacks[0].append(TreeEmbedding(current_num, True))
+#                     if len(current_embeddings_stacks[0]) > 0 and current_embeddings_stacks[0][-1].terminal:
+#                         current_left_childs.append(current_embeddings_stacks[0][-1].embedding)
+#                     else:
+#                         current_left_childs.append(None)
+#                     current_beams.append(TreeBeam(b.score + float(tv), current_node_stack, current_embeddings_stacks, current_left_childs, current_out))
+#             beams = sorted(current_beams, key=lambda x: x.score, reverse=True)
+#             beams = beams[:beam_size]
+#             flag = True
+#             for b in beams:
+#                 if len(b.node_stack[0]) != 0:
+#                     flag = False
+#             if flag:
+#                 break
+# 
+#         return beams[0].out
+# 
+#     def get_all_number_encoder_outputs(self, encoder_outputs, num_pos, batch_size, num_size, hidden_size):
+#         indices = list()
+#         sen_len = encoder_outputs.size(0)
+#         masked_index = []
+#         temp_1 = [1 for _ in range(hidden_size)]
+#         temp_0 = [0 for _ in range(hidden_size)]
+#         for b in range(batch_size):
+#             for i in num_pos[b]:
+#                 indices.append(i + b * sen_len)
+#                 masked_index.append(temp_0)
+#             indices += [0 for _ in range(len(num_pos[b]), num_size)]
+#             masked_index += [temp_1 for _ in range(len(num_pos[b]), num_size)]
+#         indices = torch.LongTensor(indices)
+#         masked_index = torch.BoolTensor(masked_index)
+#         masked_index = masked_index.view(batch_size, num_size, hidden_size)
+#         if self.USE_CUDA:
+#             indices = indices.cuda()
+#             masked_index = masked_index.cuda()
+#         all_outputs = encoder_outputs.transpose(0, 1).contiguous()
+#         all_embedding = all_outputs.view(-1, encoder_outputs.size(2))  # S x B x H -> (B x S) x H
+#         all_num = all_embedding.index_select(0, indices)
+#         all_num = all_num.view(batch_size, num_size, hidden_size)
+#         return all_num.masked_fill_(masked_index, 0.0)
+# 
+#     def generate_tree_input(self, target, decoder_output, nums_stack_batch, num_start, unk):
+#         # when the decoder input is copied num but the num has two pos, chose the max
+#         target_input = copy.deepcopy(target)
+#         for i in range(len(target)):
+#             if target[i] == unk:
+#                 num_stack = nums_stack_batch[i].pop()
+#                 max_score = -float("1e12")
+#                 for num in num_stack:
+#                     if decoder_output[i, num_start + num] > max_score:
+#                         target[i] = num + num_start
+#                         max_score = decoder_output[i, num_start + num]
+#             if target_input[i] >= num_start:
+#                 target_input[i] = 0
+#         return torch.LongTensor(target), torch.LongTensor(target_input)
+# 
+#     def mse_loss(self, outputs, targets, mask=None):
+#         # outputs   : [batch_size,output_len,hidden_size]
+#         # targets   : [batch_size,output_len,hidden_size]
+#         # mask      : [batch_size,output_len]
+#         mask = mask.to(self.device)
+#         x = torch.sqrt(torch.sum(torch.square((outputs - targets)), dim=-1))  # [batch_size,output_len]
+#         y = torch.sum(x * mask, dim=-1) / torch.sum(mask, dim=-1)  # [batch_size]
+#         return torch.sum(y)
+# 
+#     def convert_idx2symbol(self, output, num_list, num_stack):
+#         #batch_size=output.size(0)
+#         '''batch_size=1'''
+#         seq_len = len(output)
+#         num_len = len(num_list)
+#         output_list = []
+#         res = []
+#         for s_i in range(seq_len):
+#             idx = output[s_i]
+#             if idx in [self.out_sos_token, self.out_eos_token, self.out_pad_token]:
+#                 break
+#             symbol = self.out_idx2symbol[idx]
+#             if "NUM" in symbol:
+#                 num_idx = self.mask_list.index(symbol)
+#                 if num_idx >= num_len:
+#                     res = []
+#                     break
+#                 res.append(num_list[num_idx])
+#             elif symbol == SpecialTokens.UNK_TOKEN:
+#                 try:
+#                     pos_list = num_stack.pop()
+#                     c = num_list[pos_list[0]]
+#                     res.append(c)
+#                 except:
+#                     return None
+#             else:
+#                 res.append(symbol)
+#         output_list.append(res)
+#         return output_list
+
+
 class SAUSolver_(nn.Module):
     def __init__(self, config, dataset):
         super().__init__()
diff --git a/mwptoolkit/model/__init__.py b/mwptoolkit/model/__init__.py
index b9d30f3..cc122f8 100644
--- a/mwptoolkit/model/__init__.py
+++ b/mwptoolkit/model/__init__.py
@@ -2,4 +2,4 @@
 from __future__ import print_function
 from __future__ import division
 
-from mwptoolkit.model import Graph2Tree,Seq2Tree,Seq2Seq,PreTrain
\ No newline at end of file
+from mwptoolkit.model import Graph2Tree,Seq2Tree,Seq2Seq,PreTrain
diff --git a/mwptoolkit/module/Decoder/tree_decoder.py b/mwptoolkit/module/Decoder/tree_decoder.py
index d2bcdd6..c8ea41e 100644
--- a/mwptoolkit/module/Decoder/tree_decoder.py
+++ b/mwptoolkit/module/Decoder/tree_decoder.py
@@ -106,6 +106,7 @@ class SARTreeDecoder(nn.Module):
     r'''
     Seq2tree decoder with Semantically-Aligned Regularization
     '''
+
     def __init__(self, hidden_size, op_nums, generate_size, dropout=0.5):
         super(SARTreeDecoder, self).__init__()
 
@@ -129,10 +130,13 @@ def __init__(self, hidden_size, op_nums, generate_size, dropout=0.5):
 
         self.attn = TreeAttention(hidden_size, hidden_size)
         self.score = Score(hidden_size * 2, hidden_size)
+        
+        self.saligned_attn = TreeAttention(hidden_size,hidden_size)
+        self.encoder_linear1 = nn.Linear(hidden_size, hidden_size)
+        self.encoder_linear2 = nn.Linear(hidden_size, hidden_size)
 
-        self.R1=nn.Linear(hidden_size,hidden_size)
-        self.R2=nn.Linear(hidden_size,hidden_size)
-
+        self.decoder_linear1 = nn.Linear(hidden_size, hidden_size)
+        self.decoder_linear2 = nn.Linear(hidden_size, hidden_size)
 
     def forward(self, node_stacks, left_childs, encoder_outputs, num_pades, padding_hidden, seq_mask, nums_mask):
         """
@@ -144,7 +148,7 @@ def forward(self, node_stacks, left_childs, encoder_outputs, num_pades, padding_
             padding_hidden (torch.Tensor): padding hidden, shape [1,hidden_size].
             seq_mask (torch.BoolTensor): sequence mask, shape [batch_size, sequence_length].
             mask_nums (torch.BoolTensor): number mask, shape [batch_size, number_size]
-        
+
         Returns:
             tuple(torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
                 num_score, number score, shape [batch_size, number_size].
@@ -174,17 +178,17 @@ def forward(self, node_stacks, left_childs, encoder_outputs, num_pades, padding_
                 c = self.dropout(c)
                 g = torch.tanh(self.concat_r(torch.cat((ld, c), 1)))
                 t = torch.sigmoid(self.concat_rg(torch.cat((ld, c), 1)))
-                node_emb=g * t
+                node_emb = g * t
                 current_node_temp.append(node_emb)
 
         current_node = torch.stack(current_node_temp)
-        #sub_tree_emb = torch.stack(sub_tree_emb)
+        # sub_tree_emb = torch.stack(sub_tree_emb)
 
         current_embeddings = self.dropout(current_node)
 
         current_attn = self.attn(current_embeddings.transpose(0, 1), encoder_outputs, seq_mask)
         current_context = current_attn.bmm(encoder_outputs.transpose(0, 1))  # B x 1 x N
-        #s_aligned_vector=self.attn(current_embeddings, encoder_outputs, seq_mask)
+        # s_aligned_vector=self.attn(current_embeddings, encoder_outputs, seq_mask)
 
         # the information to get the current quantity
         batch_size = current_embeddings.size(0)
@@ -210,7 +214,7 @@ def forward(self, node_stacks, left_childs, encoder_outputs, num_pades, padding_
 
         return num_score, op, current_node, current_context, embedding_weight
 
-    def Semantically_Aligned_Regularization(self,subtree_emb, s_aligned_vector):
+    def Semantically_Aligned_Regularization(self, subtree_emb, s_aligned_vector):
         """
         Args:
             subtree_emb (torch.Tensor):
@@ -221,9 +225,135 @@ def Semantically_Aligned_Regularization(self,subtree_emb, s_aligned_vector):
                 s_aligned_a
                 s_aligned_d
         """
-        s_aligned_a=self.R2(torch.tanh(self.R1(s_aligned_vector)))
-        s_aligned_d=self.R2(torch.tanh(self.R1(subtree_emb)))
-        return s_aligned_a,s_aligned_d
+        #s_aligned_a = self.R2(torch.tanh(self.R1(s_aligned_vector)))
+        #s_aligned_d = self.R2(torch.tanh(self.R1(subtree_emb)))
+        s_aligned_a = self.encoder_linear2(torch.tanh(self.encoder_linear1(s_aligned_vector)))
+        s_aligned_d = self.decoder_linear2(torch.tanh(self.decoder_linear1(subtree_emb)))
+        return s_aligned_a, s_aligned_d
+
+
+# class SARTreeDecoder(nn.Module):
+#     r'''
+#     Seq2tree decoder with Semantically-Aligned Regularization
+#     '''
+#     def __init__(self, hidden_size, op_nums, generate_size, dropout=0.5):
+#         super(SARTreeDecoder, self).__init__()
+# 
+#         # Keep for reference
+#         self.hidden_size = hidden_size
+#         self.generate_size = generate_size
+#         self.op_nums = op_nums
+# 
+#         # Define layers
+#         self.dropout = nn.Dropout(dropout)
+# 
+#         self.embedding_weight = nn.Parameter(torch.randn(1, generate_size, hidden_size))
+# 
+#         # for Computational symbols and Generated numbers
+#         self.concat_l = nn.Linear(hidden_size, hidden_size)
+#         self.concat_r = nn.Linear(hidden_size * 2, hidden_size)
+#         self.concat_lg = nn.Linear(hidden_size, hidden_size)
+#         self.concat_rg = nn.Linear(hidden_size * 2, hidden_size)
+# 
+#         self.ops = nn.Linear(hidden_size * 2, op_nums)
+# 
+#         self.attn = TreeAttention(hidden_size, hidden_size)
+#         self.score = Score(hidden_size * 2, hidden_size)
+# 
+#         self.R1=nn.Linear(hidden_size,hidden_size)
+#         self.R2=nn.Linear(hidden_size,hidden_size)
+# 
+# 
+#     def forward(self, node_stacks, left_childs, encoder_outputs, num_pades, padding_hidden, seq_mask, nums_mask):
+#         """
+#         Args:
+#             node_stacks (list): node stacks.
+#             left_childs (list): representation of left childs.
+#             encoder_outputs (torch.Tensor): output from encoder, shape [sequence_length, batch_size, hidden_size].
+#             num_pades (torch.Tensor): number representation, shape [batch_size, number_size, hidden_size].
+#             padding_hidden (torch.Tensor): padding hidden, shape [1,hidden_size].
+#             seq_mask (torch.BoolTensor): sequence mask, shape [batch_size, sequence_length].
+#             mask_nums (torch.BoolTensor): number mask, shape [batch_size, number_size]
+#         
+#         Returns:
+#             tuple(torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
+#                 num_score, number score, shape [batch_size, number_size].
+#                 op, operator score, shape [batch_size, operator_size].
+#                 current_node, current node representation, shape [batch_size, 1, hidden_size].
+#                 current_context, current context representation, shape [batch_size, 1, hidden_size].
+#                 embedding_weight, embedding weight, shape [batch_size, number_size, hidden_size].
+#         """
+#         current_embeddings = []
+# 
+#         for st in node_stacks:
+#             if len(st) == 0:
+#                 current_embeddings.append(padding_hidden)
+#             else:
+#                 current_node = st[-1]
+#                 current_embeddings.append(current_node.embedding)
+# 
+#         current_node_temp = []
+#         for l, c in zip(left_childs, current_embeddings):
+#             if l is None:
+#                 c = self.dropout(c)
+#                 g = torch.tanh(self.concat_l(c))
+#                 t = torch.sigmoid(self.concat_lg(c))
+#                 current_node_temp.append(g * t)
+#             else:
+#                 ld = self.dropout(l)
+#                 c = self.dropout(c)
+#                 g = torch.tanh(self.concat_r(torch.cat((ld, c), 1)))
+#                 t = torch.sigmoid(self.concat_rg(torch.cat((ld, c), 1)))
+#                 node_emb=g * t
+#                 current_node_temp.append(node_emb)
+# 
+#         current_node = torch.stack(current_node_temp)
+#         #sub_tree_emb = torch.stack(sub_tree_emb)
+# 
+#         current_embeddings = self.dropout(current_node)
+# 
+#         current_attn = self.attn(current_embeddings.transpose(0, 1), encoder_outputs, seq_mask)
+#         current_context = current_attn.bmm(encoder_outputs.transpose(0, 1))  # B x 1 x N
+#         #s_aligned_vector=self.attn(current_embeddings, encoder_outputs, seq_mask)
+# 
+#         # the information to get the current quantity
+#         batch_size = current_embeddings.size(0)
+#         # predict the output (this node corresponding to output(number or operator)) with PADE
+# 
+#         repeat_dims = [1] * self.embedding_weight.dim()
+#         repeat_dims[0] = batch_size
+#         embedding_weight = self.embedding_weight.repeat(*repeat_dims)  # B x input_size x N
+#         embedding_weight = torch.cat((embedding_weight, num_pades), dim=1)  # B x O x N
+# 
+#         leaf_input = torch.cat((current_node, current_context), 2)
+#         leaf_input = leaf_input.squeeze(1)
+#         leaf_input = self.dropout(leaf_input)
+# 
+#         # p_leaf = nn.functional.softmax(self.is_leaf(leaf_input), 1)
+#         # max pooling the embedding_weight        
+#         embedding_weight_ = self.dropout(embedding_weight)
+#         num_score = self.score(leaf_input.unsqueeze(1), embedding_weight_, nums_mask)
+# 
+#         # num_score = nn.functional.softmax(num_score, 1)
+# 
+#         op = self.ops(leaf_input)
+# 
+#         return num_score, op, current_node, current_context, embedding_weight
+# 
+#     def Semantically_Aligned_Regularization(self,subtree_emb, s_aligned_vector):
+#         """
+#         Args:
+#             subtree_emb (torch.Tensor):
+#             s_aligned_vector (torch.Tensor):
+# 
+#         Returns:
+#             tuple(torch.Tensor, torch.Tensor):
+#                 s_aligned_a
+#                 s_aligned_d
+#         """
+#         s_aligned_a=self.R2(torch.tanh(self.R1(s_aligned_vector)))
+#         s_aligned_d=self.R2(torch.tanh(self.R1(subtree_emb)))
+#         return s_aligned_a,s_aligned_d
 
 
 class RNNBasedTreeDecoder(nn.Module):
diff --git a/mwptoolkit/module/Layer/tree_layers.py b/mwptoolkit/module/Layer/tree_layers.py
index 7326c16..42cdb57 100644
--- a/mwptoolkit/module/Layer/tree_layers.py
+++ b/mwptoolkit/module/Layer/tree_layers.py
@@ -638,4 +638,38 @@ def forward(self, hidden, encoder_outputs, seq_mask=None):
             attn_energies = attn_energies.masked_fill_(seq_mask.bool(), -1e12)
         attn_energies = nn.functional.softmax(attn_energies, dim=1)  # B x S
 
-        return attn_energies.unsqueeze(1)
\ No newline at end of file
+        return attn_energies.unsqueeze(1)
+    
+
+class SemanticAlignmentModule(nn.Module):
+    def __init__(self, encoder_hidden_size, decoder_hidden_size, hidden_size, batch_first=False):
+        super(SemanticAlignmentModule, self).__init__()
+        self.batch_first = batch_first
+        self.attn = TreeAttention(encoder_hidden_size,decoder_hidden_size)
+        self.encoder_linear1 = nn.Linear(encoder_hidden_size, hidden_size)
+        self.encoder_linear2 = nn.Linear(hidden_size, hidden_size)
+
+        self.decoder_linear1 = nn.Linear(decoder_hidden_size, hidden_size)
+        self.decoder_linear2 = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, decoder_hidden, encoder_outputs):
+        if self.batch_first:
+            decoder_hidden = decoder_hidden.unsqueeze(0)
+            encoder_outputs = encoder_outputs.unsqueeze(0)
+        else:
+            decoder_hidden = decoder_hidden.unsqueeze(0)
+            encoder_outputs = encoder_outputs.unsqueeze(1)
+        attn_weights = self.attn(decoder_hidden, encoder_outputs, None)
+        if self.batch_first:
+            align_context = attn_weights.bmm(encoder_outputs) # B x 1 x H
+        else:
+            align_context = attn_weights.bmm(encoder_outputs.transpose(0, 1))  # B x 1 x H
+            align_context = align_context.transpose(0,1)
+
+        encoder_linear1 = torch.tanh(self.encoder_linear1(align_context))
+        encoder_linear2 = self.encoder_linear2(encoder_linear1)
+
+        decoder_linear1 = torch.tanh(self.decoder_linear1(decoder_hidden))
+        decoder_linear2 = self.decoder_linear2(decoder_linear1)
+
+        return encoder_linear2, decoder_linear2
diff --git a/mwptoolkit/trainer/supervised_trainer.py b/mwptoolkit/trainer/supervised_trainer.py
index 9ee7a39..6940865 100644
--- a/mwptoolkit/trainer/supervised_trainer.py
+++ b/mwptoolkit/trainer/supervised_trainer.py
@@ -1107,23 +1107,24 @@ def param_search(self):
 
 class SAUSolverTrainer(GTSTrainer):
     """sausolver trainer, used to implement training, testing, parameter searching for deep-learning model SAUSolver.
-    
+
     example of instantiation:
-        
+
         >>> trainer = SAUSolverTrainer(config, model, dataloader, evaluator)
 
         for training:
-            
+
             >>> trainer.fit()
-        
+
         for testing:
-            
+
             >>> trainer.test()
-        
+
         for parameter searching:
 
             >>> trainer.param_search()
     """
+
     def __init__(self, config, model, dataloader, evaluator):
         """
         Args:
@@ -1131,7 +1132,7 @@ def __init__(self, config, model, dataloader, evaluator):
             model (Model): An object of deep-learning model. 
             dataloader (Dataloader): dataloader object.
             evaluator (Evaluator): evaluator object.
-        
+
         expected that config includes these parameters below:
 
         learning_rate (float): learning rate of model.
@@ -1163,15 +1164,12 @@ def _train_batch(self, batch):
         return batch_loss
 
     def _eval_batch(self, batch):
-        try:
-            test_out, target = self.model.model_test(batch)
-        except:
-            print(batch['id'])
-
+        test_out, target = self.model.model_test(batch)
         batch_size = len(test_out)
         val_acc = []
         equ_acc = []
         for idx in range(batch_size):
+            # batch['ans'][idx] = [12,8]
             if self.config["task_type"] == TaskType.SingleEquation:
                 val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx])
             elif self.config["task_type"] == TaskType.MultiEquation:
@@ -1180,18 +1178,114 @@ def _eval_batch(self, batch):
                 raise NotImplementedError
             val_acc.append(val_ac)
             equ_acc.append(equ_ac)
-            result={
-                'id':batch['id'][idx],
-                'prediction':' '.join(test_out[idx]),
-                'target':' '.join(target[idx]),
-                'number list':batch['num list'][idx],
-                'value acc':val_ac,
-                'equ acc':equ_ac
+            result = {
+                'id': batch['id'][idx],
+                'prediction': ' '.join(test_out[idx]),
+                'target': ' '.join(target[idx]),
+                'number list': batch['num list'][idx],
+                'value acc': val_ac,
+                'equ acc': equ_ac
             }
             self.output_result.append(result)
         return val_acc, equ_acc
 
+    def _build_optimizer(self):
+        self.embedder_optimizer = torch.optim.Adam(self.model.embedder.parameters(),
+                                                   self.config["embedding_learning_rate"],
+                                                   weight_decay=self.config["weight_decay"])
+        self.encoder_optimizer = torch.optim.Adam(self.model.encoder.parameters(), self.config["learning_rate"],
+                                                  weight_decay=self.config["weight_decay"])
+        self.decoder_optimizer = torch.optim.Adam(self.model.decoder.parameters(), self.config["learning_rate"],
+                                                  weight_decay=self.config["weight_decay"])
+        self.node_generater_optimizer = torch.optim.Adam(self.model.node_generater.parameters(),
+                                                         self.config["learning_rate"],
+                                                         weight_decay=self.config["weight_decay"])
+        self.merge_optimizer = torch.optim.Adam(self.model.merge.parameters(), self.config["learning_rate"],
+                                                weight_decay=self.config["weight_decay"])
+        self.sa_optimizer = torch.optim.Adam(self.model.sa.parameters(), self.config['learning_rate'],
+                                             weight_decay=self.config["weight_decay"])
+        # scheduler
+        self.embedder_scheduler = torch.optim.lr_scheduler.StepLR(self.embedder_optimizer,
+                                                                  step_size=self.config["step_size"], gamma=0.5, )
+        self.encoder_scheduler = torch.optim.lr_scheduler.StepLR(self.encoder_optimizer,
+                                                                 step_size=self.config["step_size"], gamma=0.5)
+        self.decoder_scheduler = torch.optim.lr_scheduler.StepLR(self.decoder_optimizer,
+                                                                 step_size=self.config["step_size"], gamma=0.5)
+        self.node_generater_scheduler = torch.optim.lr_scheduler.StepLR(self.node_generater_optimizer,
+                                                                        step_size=self.config["step_size"], gamma=0.5)
+        self.merge_scheduler = torch.optim.lr_scheduler.StepLR(self.merge_optimizer, step_size=self.config["step_size"],
+                                                               gamma=0.5)
+        self.sa_scheduler = torch.optim.lr_scheduler.StepLR(self.sa_optimizer, step_size=self.config['step_size'],
+                                                            gamma=0.5)
+
+    def _save_checkpoint(self):
+        check_pnt = {
+            "model": self.model.state_dict(),
+            "embedder_optimizer": self.embedder_optimizer.state_dict(),
+            "encoder_optimizer": self.encoder_optimizer.state_dict(),
+            "decoder_optimizer": self.decoder_optimizer.state_dict(),
+            "generate_optimizer": self.node_generater_optimizer.state_dict(),
+            "merge_optimizer": self.merge_optimizer.state_dict(),
+            "sa_optimizer":self.sa_optimizer.state_dict(),
+            "embedder_scheduler": self.embedder_scheduler.state_dict(),
+            "encoder_scheduler": self.encoder_scheduler.state_dict(),
+            "decoder_scheduler": self.decoder_scheduler.state_dict(),
+            "generate_scheduler": self.node_generater_scheduler.state_dict(),
+            "merge_scheduler": self.merge_scheduler.state_dict(),
+            "sa_scheduler":self.sa_scheduler.state_dict(),
+            "start_epoch": self.epoch_i,
+            "best_valid_value_accuracy": self.best_valid_value_accuracy,
+            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
+            "best_test_value_accuracy": self.best_test_value_accuracy,
+            "best_test_equ_accuracy": self.best_test_equ_accuracy,
+            "best_folds_accuracy": self.best_folds_accuracy,
+            "fold_t": self.config["fold_t"]
+        }
+        torch.save(check_pnt, self.config["checkpoint_path"])
+
+    def _load_checkpoint(self):
+        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
+        # load parameter of model
+        self.model.load_state_dict(check_pnt["model"])
+        # load parameter of optimizer
+        self.embedder_optimizer.load_state_dict(check_pnt["embedder_optimizer"])
+        self.encoder_optimizer.load_state_dict(check_pnt["encoder_optimizer"])
+        self.decoder_optimizer.load_state_dict(check_pnt["decoder_optimizer"])
+        self.node_generater_optimizer.load_state_dict(check_pnt["generate_optimizer"])
+        self.merge_optimizer.load_state_dict(check_pnt["merge_optimizer"])
+        self.sa_optimizer.load_state_dict(check_pnt["sa_optimizer"])
+        # load parameter of scheduler
+        self.embedder_scheduler.load_state_dict(check_pnt['embedder_scheduler'])
+        self.encoder_scheduler.load_state_dict(check_pnt["encoder_scheduler"])
+        self.decoder_scheduler.load_state_dict(check_pnt["decoder_scheduler"])
+        self.node_generater_scheduler.load_state_dict(check_pnt["generate_scheduler"])
+        self.merge_scheduler.load_state_dict(check_pnt["merge_scheduler"])
+        self.sa_scheduler.load_state_dict(check_pnt["sa_scheduler"])
+        # other parameter
+        self.start_epoch = check_pnt["start_epoch"]
+        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
+        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
+        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
+        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
+        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
+
+    def _scheduler_step(self):
+        self.embedder_scheduler.step()
+        self.encoder_scheduler.step()
+        self.decoder_scheduler.step()
+        self.node_generater_scheduler.step()
+        self.merge_scheduler.step()
+        self.sa_scheduler.step()
+
+    def _optimizer_step(self):
+        self.embedder_optimizer.step()
+        self.encoder_optimizer.step()
+        self.decoder_optimizer.step()
+        self.node_generater_optimizer.step()
+        self.merge_optimizer.step()
+        self.sa_optimizer.step()
 
+    
 class TRNNTrainer(SupervisedTrainer):
     """trnn trainer, used to implement training, testing, parameter searching for deep-learning model TRNN.
     
@@ -1492,6 +1586,7 @@ def param_search(self):
 
                 tune.report(accuracy=test_val_ac)
 
+
 class SalignedTrainer(SupervisedTrainer):
     """saligned trainer, used to implement training, testing, parameter searching for deep-learning model S-aligned.
     
@@ -1980,7 +2075,49 @@ def _save_checkpoint(self):
         torch.save(check_pnt, self.config["checkpoint_path"])
 
     def _load_checkpoint(self):
-        return super()._load_checkpoint()
+        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
+        # load parameter of model
+        self.model.load_state_dict(check_pnt["model"])
+        # load parameter of optimizer
+        self.t_embedder_optimizer.load_state_dict(check_pnt['t_embedder_optimizer'])
+        self.t_encoder_optimizer.load_state_dict(check_pnt['t_encoder_optimizer'])
+        self.t_decoder_optimizer.load_state_dict(check_pnt['t_decoder_optimizer'])
+        self.t_node_generater_optimizer.load_state_dict(check_pnt['t_node_generater_optimizer'])
+        self.t_merge_optimizer.load_state_dict(check_pnt['t_merge_optimizer'])
+
+        self.s_embedder_optimizer.load_state_dict(check_pnt['s_embedder_optimizer'])
+        self.s_encoder_optimizer.load_state_dict(check_pnt['s_encoder_optimizer'])
+        self.s_decoder_optimizer1.load_state_dict(check_pnt['s_decoder_optimizer1'])
+        self.s_node_generater_optimizer1.load_state_dict(check_pnt['s_node_generater_optimizer1'])
+        self.s_merge_optimizer1.load_state_dict(check_pnt['s_merge_optimizer1'])
+        self.s_decoder_optimizer2.load_state_dict(check_pnt['s_decoder_optimizer2'])
+        self.s_node_generater_optimizer2.load_state_dict(check_pnt['s_node_generater_optimizer2'])
+        self.s_merge_optimizer2.load_state_dict(check_pnt['s_merge_optimizer2'])
+
+        # load parameter of scheduler
+        self.t_embedder_scheduler.load_state_dict(check_pnt['t_embedder_scheduler'])
+        self.t_encoder_scheduler.load_state_dict(check_pnt['t_encoder_scheduler'])
+        self.t_decoder_scheduler.load_state_dict(check_pnt['t_decoder_scheduler'])
+        self.t_node_generater_scheduler.load_state_dict(check_pnt['t_node_generater_scheduler'])
+        self.t_merge_scheduler.load_state_dict(check_pnt['t_merge_scheduler'])
+
+        self.s_embedder_scheduler.load_state_dict(check_pnt['s_embedder_scheduler'])
+        self.s_encoder_scheduler.load_state_dict(check_pnt['s_encoder_scheduler'])
+        self.s_decoder_scheduler1.load_state_dict(check_pnt['s_decoder_scheduler1'])
+        self.s_node_generater_scheduler1.load_state_dict(check_pnt['s_node_generater_scheduler1'])
+        self.s_merge_scheduler1.load_state_dict(check_pnt['s_merge_scheduler1'])
+        self.s_decoder_scheduler2.load_state_dict(check_pnt['s_decoder_scheduler2'])
+        self.s_node_generater_scheduler2.load_state_dict(check_pnt['s_node_generater_scheduler2'])
+        self.s_merge_scheduler2.load_state_dict(check_pnt['s_merge_scheduler2'])
+
+        # other parameter
+        self.t_start_epoch = check_pnt["t_start_epoch"]
+        self.s_start_epoch = check_pnt['s_start_epoch']
+        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
+        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
+        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
+        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
+        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
 
     def _teacher_net_train(self):
         self.model.t_embedder.train()
@@ -2736,4 +2873,94 @@ def _build_optimizer(self):
             ],
             lr = self.config["ans_learning_rate"],
             momentum = 0.9
-        )
\ No newline at end of file
+        )
+
+class MWPBertTrainer(GTSTrainer):
+    def __init__(self,config,model,dataloader,evaluator):
+        super().__init__(config,model,dataloader,evaluator)
+    
+    def _build_optimizer(self):
+        self.encoder_optimizer = torch.optim.Adam(
+            self.model.encoder.parameters(), 
+            self.config['encoding_learning_rate'],
+            weight_decay=self.config["weight_decay"]
+        )
+        self.decoder_optimizer = torch.optim.Adam(
+            self.model.decoder.parameters(), 
+            self.config["learning_rate"],
+            weight_decay=self.config["weight_decay"]
+        )
+        self.node_generater_optimizer = torch.optim.Adam(
+            self.model.node_generater.parameters(),
+            self.config["learning_rate"],
+            weight_decay=self.config["weight_decay"]
+        )
+        self.merge_optimizer = torch.optim.Adam(
+            self.model.merge.parameters(), 
+            self.config["learning_rate"],
+            weight_decay=self.config["weight_decay"]
+        )
+        # scheduler
+        self.encoder_scheduler = torch.optim.lr_scheduler.StepLR(self.encoder_optimizer,
+                                                                 step_size=self.config["step_size"], gamma=0.5)
+        self.decoder_scheduler = torch.optim.lr_scheduler.StepLR(self.decoder_optimizer,
+                                                                 step_size=self.config["step_size"], gamma=0.5)
+        self.node_generater_scheduler = torch.optim.lr_scheduler.StepLR(self.node_generater_optimizer,
+                                                                        step_size=self.config["step_size"], gamma=0.5)
+        self.merge_scheduler = torch.optim.lr_scheduler.StepLR(self.merge_optimizer, step_size=self.config["step_size"],
+                                                               gamma=0.5)
+    
+    def _save_checkpoint(self):
+        check_pnt = {
+            "model": self.model.state_dict(),
+            "encoder_optimizer": self.encoder_optimizer.state_dict(),
+            "decoder_optimizer": self.decoder_optimizer.state_dict(),
+            "generate_optimizer": self.node_generater_optimizer.state_dict(),
+            "merge_optimizer": self.merge_optimizer.state_dict(),
+            "encoder_scheduler": self.encoder_scheduler.state_dict(),
+            "decoder_scheduler": self.decoder_scheduler.state_dict(),
+            "generate_scheduler": self.node_generater_scheduler.state_dict(),
+            "merge_scheduler": self.merge_scheduler.state_dict(),
+            "start_epoch": self.epoch_i,
+            "best_valid_value_accuracy": self.best_valid_value_accuracy,
+            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
+            "best_test_value_accuracy": self.best_test_value_accuracy,
+            "best_test_equ_accuracy": self.best_test_equ_accuracy,
+            "best_folds_accuracy": self.best_folds_accuracy,
+            "fold_t": self.config["fold_t"]
+        }
+        torch.save(check_pnt, self.config["checkpoint_path"])
+
+    def _load_checkpoint(self):
+        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
+        # load parameter of model
+        self.model.load_state_dict(check_pnt["model"])
+        # load parameter of optimizer
+        self.encoder_optimizer.load_state_dict(check_pnt["encoder_optimizer"])
+        self.decoder_optimizer.load_state_dict(check_pnt["decoder_optimizer"])
+        self.node_generater_optimizer.load_state_dict(check_pnt["generate_optimizer"])
+        self.merge_optimizer.load_state_dict(check_pnt["merge_optimizer"])
+        # load parameter of scheduler
+        self.encoder_scheduler.load_state_dict(check_pnt["encoder_scheduler"])
+        self.decoder_scheduler.load_state_dict(check_pnt["decoder_scheduler"])
+        self.node_generater_scheduler.load_state_dict(check_pnt["generate_scheduler"])
+        self.merge_scheduler.load_state_dict(check_pnt["merge_scheduler"])
+        # other parameter
+        self.start_epoch = check_pnt["start_epoch"]
+        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
+        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
+        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
+        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
+        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
+
+    def _scheduler_step(self):
+        self.encoder_scheduler.step()
+        self.decoder_scheduler.step()
+        self.node_generater_scheduler.step()
+        self.merge_scheduler.step()
+
+    def _optimizer_step(self):
+        self.encoder_optimizer.step()
+        self.decoder_optimizer.step()
+        self.node_generater_optimizer.step()
+        self.merge_optimizer.step()
diff --git a/mwptoolkit/trainer/template_trainer.py b/mwptoolkit/trainer/template_trainer.py
index b47f7b1..ddf41d1 100644
--- a/mwptoolkit/trainer/template_trainer.py
+++ b/mwptoolkit/trainer/template_trainer.py
@@ -6,10 +6,11 @@
 
 import time
 
-from mwptoolkit.trainer.trainer import AbstractTrainer
+from mwptoolkit.trainer.abstract_trainer import AbstractTrainer
 from mwptoolkit.utils.enum_type import DatasetType
 from mwptoolkit.utils.utils import time_since
 
+
 class TemplateTrainer(AbstractTrainer):
     r"""template trainer.
     
diff --git a/mwptoolkit/trainer/trainer.py b/mwptoolkit/trainer/trainer.py
deleted file mode 100644
index e263b66..0000000
--- a/mwptoolkit/trainer/trainer.py
+++ /dev/null
@@ -1,2739 +0,0 @@
-import torch
-import time
-from logging import getLogger
-from torch import nn
-import numpy as np
-from mwptoolkit.utils.utils import time_since
-from mwptoolkit.utils.enum_type import PAD_TOKEN, DatasetType,TaskType,SpecialTokens
-from mwptoolkit.loss.masked_cross_entropy_loss import MaskedCrossEntropyLoss
-from mwptoolkit.loss.nll_loss import NLLLoss
-from mwptoolkit.loss.binary_cross_entropy_loss import BinaryCrossEntropyLoss
-from mwptoolkit.module.Optimizer.optim import WarmUpScheduler
-from mwptoolkit.loss.mse_loss import MSELoss
-
-
-class AbstractTrainer(object):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__()
-        self.config = config
-        self.model = model
-        self.dataloader = dataloader
-        self.evaluator = evaluator
-        self.logger = getLogger()
-        self.best_folds_accuracy=config["best_folds_accuracy"]
-        self.test_step=config["test_step"]
-
-        self.best_valid_equ_accuracy = 0.
-        self.best_valid_value_accuracy = 0.
-        self.best_test_equ_accuracy = 0.
-        self.best_test_value_accuracy = 0.
-        self.start_epoch = 0
-        self.epoch_i = 0
-
-    def _save_checkpoint(self):
-        raise NotImplementedError
-
-    def _load_checkpoint(self):
-        raise NotImplementedError
-
-    def _save_model(self):
-        state_dict = {"model": self.model.state_dict()}
-        if self.config["k_fold"]:
-            path=self.config["trained_model_path"][:-4]+"-fold{}.pth".format(self.config["fold_t"])
-            torch.save(state_dict,path)
-        else:
-            torch.save(state_dict, self.config["trained_model_path"])
-
-    def _load_model(self):
-        if self.config["k_fold"]:
-            path=self.config["trained_model_path"][:-4]+"-fold{}.pth".format(self.config["fold_t"])
-            state_dict = torch.load(path, map_location=self.config["map_location"])
-        else:
-            state_dict = torch.load(self.config["trained_model_path"], map_location=self.config["map_location"])
-        #self.model = nn.DataParallel(self.model)
-        self.model.load_state_dict(state_dict["model"],False)
-
-    def _build_optimizer(self):
-        raise NotImplementedError
-
-    def _train_batch(self):
-        raise NotADirectoryError
-
-    def _eval_batch(self):
-        raise NotImplementedError
-
-    def _train_epoch(self):
-        raise NotImplementedError
-
-    def fit(self):
-        raise NotImplementedError
-
-    def evaluate(self, eval_set):
-        raise NotImplementedError
-
-    def test(self):
-        raise NotImplementedError
-
-
-class Trainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"], self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _build_loss(self, symbol_size, out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.loss = NLLLoss(weight, pad)
-
-    def _idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.out_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-
-    def _train_batch(self, batch):
-        outputs = self.model(batch["question"], batch["ques len"], batch["equation"])
-        #outputs=torch.nn.functional.log_softmax(outputs,dim=1)
-        if self.config["share_vocab"]:
-            batch_equation = self._idx2word_2idx(batch["equation"])
-            self.loss.eval_batch(outputs, batch_equation.view(-1))
-        else:
-            self.loss.eval_batch(outputs, batch["equation"].view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        test_out = self.model(batch["question"], batch["ques len"])
-        if self.config["share_vocab"]:
-            target = self._idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            if self.config["task_type"]==TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            elif self.config["task_type"]==TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            else:
-                raise NotImplementedError
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self.optimizer.step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                self.logger.info("----------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                self.logger.info("----------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                if valid_val_ac >= self.best_valid_value_accuracy:
-                    self.best_valid_value_accuracy = valid_val_ac
-                    self.best_valid_equ_accuracy = valid_equ_ac
-                    self.best_test_value_accuracy = test_val_ac
-                    self.best_test_equ_accuracy = test_equ_ac
-                    self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-    
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-
-class SingleEquationTrainer(Trainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"], self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _build_loss(self, symbol_size, out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.loss = NLLLoss(weight, pad)
-
-    def _idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.out_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-
-    def _train_batch(self, batch):
-        outputs = self.model(batch["question"], batch["ques len"], batch["equation"])
-        #outputs=torch.nn.functional.log_softmax(outputs,dim=1)
-        if self.config["share_vocab"]:
-            batch_equation = self._idx2word_2idx(batch["equation"])
-            self.loss.eval_batch(outputs, batch_equation.view(-1))
-        else:
-            self.loss.eval_batch(outputs, batch["equation"].view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        test_out = self.model(batch["question"], batch["ques len"])
-        if self.config["share_vocab"]:
-            target = self._idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            #val_ac, equ_ac, _, _ = self.evaluator.result(target[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self.optimizer.step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-    
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-            
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-class MultiEquationTrainer(Trainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"], self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _build_loss(self, symbol_size, out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.loss = NLLLoss(weight, pad)
-
-    def _idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.out_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-
-    def _train_batch(self, batch):
-        outputs = self.model(batch["question"], batch["ques len"], batch["equation"])
-        #outputs=torch.nn.functional.log_softmax(outputs,dim=1)
-        if self.config["share_vocab"]:
-            batch_equation = self._idx2word_2idx(batch["equation"])
-            self.loss.eval_batch(outputs, batch_equation.view(-1))
-        else:
-            self.loss.eval_batch(outputs, batch["equation"].view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        test_out = self.model(batch["question"], batch["ques len"])
-        if self.config["share_vocab"]:
-            target = self._idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            #val_ac, equ_ac, _, _ = self.evaluator.result_multi(target[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self.optimizer.step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-
-class GTSTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self.loss = MaskedCrossEntropyLoss()
-
-    def _build_optimizer(self):
-        # optimizer
-        self.embedder_optimizer = torch.optim.Adam(self.model.embedder.parameters(), self.config["learning_rate"], weight_decay=self.config["weight_decay"])
-        self.encoder_optimizer = torch.optim.Adam(self.model.encoder.parameters(), self.config["learning_rate"], weight_decay=self.config["weight_decay"])
-        self.decoder_optimizer = torch.optim.Adam(self.model.parameters(), self.config["learning_rate"], weight_decay=self.config["weight_decay"])
-        self.node_generater_optimizer = torch.optim.Adam(self.model.node_generater.parameters(), self.config["learning_rate"], weight_decay=self.config["weight_decay"])
-        self.merge_optimizer = torch.optim.Adam(self.model.merge.parameters(), self.config["learning_rate"], weight_decay=self.config["weight_decay"])
-        # scheduler
-        self.embedder_scheduler = torch.optim.lr_scheduler.StepLR(self.embedder_optimizer, step_size=self.config["step_size"], gamma=0.5)
-        self.encoder_scheduler = torch.optim.lr_scheduler.StepLR(self.encoder_optimizer, step_size=self.config["step_size"], gamma=0.5)
-        self.decoder_scheduler = torch.optim.lr_scheduler.StepLR(self.decoder_optimizer, step_size=self.config["step_size"], gamma=0.5)
-        self.node_generater_scheduler = torch.optim.lr_scheduler.StepLR(self.node_generater_optimizer, step_size=self.config["step_size"], gamma=0.5)
-        self.merge_scheduler = torch.optim.lr_scheduler.StepLR(self.merge_optimizer, step_size=self.config["step_size"], gamma=0.5)
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "embedder_optimizer": self.embedder_optimizer.state_dict(),
-            "encoder_optimizer": self.encoder_optimizer.state_dict(),
-            "decoder_optimizer": self.decoder_optimizer.state_dict(),
-            "generate_optimizer": self.node_generater_optimizer.state_dict(),
-            "merge_optimizer": self.merge_optimizer.state_dict(),
-            "embedder_scheduler": self.embedder_scheduler.state_dict(),
-            "encoder_scheduler": self.encoder_scheduler.state_dict(),
-            "decoder_scheduler": self.decoder_scheduler.state_dict(),
-            "generate_scheduler": self.node_generater_scheduler.state_dict(),
-            "merge_scheduler": self.merge_scheduler.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.embedder_optimizer.load_state_dict(check_pnt["embedder_optimizer"])
-        self.encoder_optimizer.load_state_dict(check_pnt["encoder_optimizer"])
-        self.decoder_optimizer.load_state_dict(check_pnt["decoder_optimizer"])
-        self.node_generater_optimizer.load_state_dict(check_pnt["generate_optimizer"])
-        self.merge_optimizer.load_state_dict(check_pnt["merge_optimizer"])
-        #load parameter of scheduler
-        self.embedder_scheduler.load_state_dict(check_pnt["embedder_scheduler"])
-        self.encoder_scheduler.load_state_dict(check_pnt["encoder_scheduler"])
-        self.decoder_scheduler.load_state_dict(check_pnt["decoder_scheduler"])
-        self.node_generater_scheduler.load_state_dict(check_pnt["generate_scheduler"])
-        self.merge_scheduler.load_state_dict(check_pnt["merge_scheduler"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _scheduler_step(self):
-        self.embedder_scheduler.step()
-        self.encoder_scheduler.step()
-        self.decoder_scheduler.step()
-        self.node_generater_scheduler.step()
-        self.merge_scheduler.step()
-
-    def _optimizer_step(self):
-        self.embedder_optimizer.step()
-        self.encoder_optimizer.step()
-        self.decoder_optimizer.step()
-        self.node_generater_optimizer.step()
-        self.merge_optimizer.step()
-
-    def _model_zero_grad(self):
-        self.model.embedder.zero_grad()
-        self.model.encoder.zero_grad()
-        self.model.decoder.zero_grad()
-        self.model.node_generater.zero_grad()
-        self.model.merge.zero_grad()
-
-    def _model_train(self):
-        self.model.embedder.train()
-        self.model.encoder.train()
-        self.model.decoder.train()
-        self.model.node_generater.train()
-        self.model.merge.train()
-
-    def _model_eval(self):
-        self.model.embedder.eval()
-        self.model.encoder.eval()
-        self.model.decoder.eval()
-        self.model.node_generater.eval()
-        self.model.merge.eval()
-
-    def _train_batch(self, batch):
-        '''
-        seq, seq_length, nums_stack, num_size, generate_nums, num_pos,\
-                UNK_TOKEN,num_start,target=None, target_length=None,max_length=30,beam_size=5
-        '''
-        unk = self.dataloader.out_unk_token
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in self.dataloader.dataset.generate_list]
-
-        outputs, _=self.model(batch["question"],batch["ques len"],batch["num stack"],batch["num size"],\
-                                generate_nums,batch["num pos"],num_start,batch["equation"],batch["equ len"],UNK_TOKEN=unk)
-        self.loss.eval_batch(outputs, batch["equation"], batch["equ mask"])
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        # if batch["id"]==[80313]:
-        #     print(1)
-        # else:
-        #     return [False], [False]
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in self.dataloader.dataset.generate_list]
-        test_out=self.model(batch["question"],batch["ques len"],batch["num stack"],batch["num size"],\
-                                generate_nums,batch["num pos"],num_start)
-
-        if self.config["task_type"]==TaskType.SingleEquation:
-            val_ac, equ_ac, _, _ = self.evaluator.result(test_out, batch["equation"].tolist()[0], batch["num list"][0], batch["num stack"][0])
-        elif self.config["task_type"]==TaskType.MultiEquation:
-            val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out, batch["equation"].tolist()[0], batch["num list"][0], batch["num stack"][0])
-        else:
-            raise NotImplementedError
-        return [val_ac], [equ_ac]
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self._model_train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self._model_zero_grad()
-            
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self._optimizer_step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-        #print("epoch [%2d]avr loss [%2.8f]"%(self.epoch_i,loss_total /self.batch_nums))
-        #print("epoch train time {}".format(time_since(time.time() -epoch_start_time)))
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self._scheduler_step()
-            
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-        
-
-    def evaluate(self, eval_set):
-        self._model_eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-class GTSWeakTrainer(GTSTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self.loss = MaskedCrossEntropyLoss()
-
-    def _build_buffer_batch(self):
-        self._buffer_batches = [[] for i in range(self.dataloader.trainset_nums)]
-        self._buffer_batches_exp = [[] for i in range(self.dataloader.trainset_nums)]
-
-    def _train_batch(self, batch):
-        '''
-        seq, seq_length, nums_stack, num_size, generate_nums, num_pos,\
-                UNK_TOKEN,num_start,target=None, target_length=None,max_length=30,beam_size=5
-        '''
-        unk = self.dataloader.out_unk_token
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in self.dataloader.dataset.generate_list]
-
-        outputs=self.model(batch["question"],batch["ques len"],batch["num stack"],batch["num size"],\
-                                generate_nums,batch["num pos"],num_start,batch["equation"],batch["equ len"],UNK_TOKEN=unk)
-        self.loss.eval_batch(outputs, batch["equation"], batch["equ mask"])
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _train_epoch(self):
-        epoch_start_time = time.time() #
-        loss_total = 0.  #
-        #self._model_train()
-        self.mask_flag = False #
-        self._pos = 0 #
-        self.epo_iteration = 0
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self._model_train()
-            self._model_zero_grad()
-            if self.epoch_i == 1 and self.batch_idx <= 2: #
-                self.mask_flag = True
-            buffer_batches_train = self._buffer_batches[self._pos: self._pos + len(batch["ques len"])]
-            buffer_batches_train_exp = self._buffer_batches_exp[self._pos: self._pos + len(batch["ques len"])]
-            iterations, buffer_batch_new, buffer_batch_exp, batch_loss = self._weakly_train_batch(batch, buffer_batches_train, buffer_batches_train_exp)
-            loss_total += batch_loss
-            self.epo_iteration += iterations
-            self._buffer_batches[self._pos: self._pos + len(batch["ques len"])] = buffer_batch_new
-            self._buffer_batches_exp[self._pos: self._pos + len(batch["ques len"])] = buffer_batch_exp
-            self._pos += len(batch["ques len"])
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        loss_total = loss_total if self.epo_iteration == 0 else loss_total/self.epo_iteration
-        return loss_total, epoch_time_cost
-
-    def _weakly_train_batch(self, batch, buffer_batches_train, buffer_batches_train_exp):
-        batch_size = len(batch["ques len"])
-
-        unk = self.dataloader.out_unk_token
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in
-                         self.dataloader.dataset.generate_list]
-        input_var = batch["question"]  # [batch_size, max_len]
-
-        input_var = input_var.transpose(0, 1)
-        fix_input_length, fix_target_list, \
-        fix_index, fix_target_length, \
-        buffer_batch_new, buffer_batch_exp = self.model.weakly_train(batch["question"], batch["ques len"], batch['ans'],
-                                                                     batch["num list"], batch["num size"], \
-                                                                     generate_nums, batch["num pos"], num_start,
-                                                                     batch["equation"], batch["equ len"],
-                                                                     self.epoch_i - 1, self.mask_flag, unk, \
-                                                                     self.config['supervising_mode'],
-                                                                     buffer_batches_train, buffer_batches_train_exp,
-                                                                     Lang=self.dataloader, n_step=50)
-        # buffer_batch_new, buffer_batch_exp = Weakly_Supervising(generate_exps_list, all_node_outputs_mask_list,buffer_batches_train, buffer_batches_train_exp)
-
-        # print("fix_input_length", fix_input_length)
-        # print("fix_target_list", fix_target_list)
-
-        fix_input_length = np.array(fix_input_length)
-        fix_target_list = np.array(fix_target_list)
-        fix_index = np.array(fix_index)
-        fix_target_length = np.array(fix_target_length)
-
-        inds = np.argsort(-fix_input_length)
-        # print("inds", inds)
-        fix_target_list = fix_target_list[inds].tolist()
-        # print("fix_target_list", fix_target_list)
-        fix_index = fix_index[inds].tolist()
-        # print("fix_index", fix_index)
-        fix_target_length = fix_target_length[inds].tolist()
-        # print("fix_target_length", fix_target_length)
-        # print(batch_size)
-        mapo_batch_size = 64
-        if not len(fix_target_list) % mapo_batch_size == 0:
-            num_iteration = int(len(fix_target_list) / mapo_batch_size) + 1
-        else:
-            num_iteration = int(len(fix_target_list) / mapo_batch_size)
-        # print("num_iteration", num_iteration)
-        batch_loss = torch.FloatTensor([[0.0]]).to(self.config["device"])
-        for j in range(num_iteration):
-            if not j * mapo_batch_size + mapo_batch_size - 1 < len(fix_target_list):
-                mapo_batch_size = len(fix_target_list) - ((j - 1) * mapo_batch_size + mapo_batch_size)
-            target_length_mapo = fix_target_length[j * mapo_batch_size: (j * mapo_batch_size + mapo_batch_size)]
-            idx_list = fix_index[j * mapo_batch_size: (j * mapo_batch_size + mapo_batch_size)]
-            target_list = fix_target_list[j * mapo_batch_size: (j * mapo_batch_size + mapo_batch_size)]
-            input_length_mapo = []
-            num_pos_mapo = []
-            num_size_mapo_batch = []
-            # print("input_length", batch["ques len"])
-
-            for k in range(mapo_batch_size):
-                idx = idx_list[k]
-                input_length_mapo.append(batch["ques len"][idx])
-                num_pos_mapo.append(batch["num pos"][idx])
-                num_size_mapo_batch.append(batch["num size"][idx])
-            # print("input_length_mapo", input_length_mapo)
-
-            input_var_mapo = torch.zeros((max(input_length_mapo), mapo_batch_size), dtype=torch.long)
-            # print("input_var_mapo", input_var_mapo.size())
-            target = torch.zeros((mapo_batch_size, max(target_length_mapo)), dtype=torch.long).to(self.config["device"])
-
-            for k in range(mapo_batch_size):
-                idx = idx_list[k]
-                input_var_mapo[:, k] = input_var[:, idx][:max(input_length_mapo)]
-                target[k][:target_length_mapo[k]] = torch.LongTensor(target_list[k])
-
-            # print("target", target.size())
-            self._model_train()
-            self._model_zero_grad()
-            input_var_mapo = input_var_mapo.to(self.config["device"])
-            input_var_mapo = input_var_mapo.transpose(0, 1)
-
-            output, target = self.model(input_var_mapo, input_length_mapo, batch["num stack"], num_size_mapo_batch,
-                                generate_nums, num_pos_mapo, \
-                                num_start, target=target, target_length=target_length_mapo, UNK_TOKEN=unk)
-            target_length_mapo = torch.LongTensor(target_length_mapo).to(self.config["device"])
-            target_mask_mapo = self.sequence_mask(target_length_mapo)
-            self.loss.eval_batch(output, target, target_mask_mapo)
-            batch_loss += self.loss.get_loss()
-            self.loss.backward()
-            self._optimizer_step()
-            self.loss.reset()
-        # self._buffer_batches[self._pos: self._pos + len(batch["ques len"])] = buffer_batch_new
-        # self._buffer_batches_exp[self._pos: self._pos + len(batch["ques len"])] = buffer_batch_exp
-
-        # self._pos += len(batch["ques len"])
-        batch_loss = batch_loss if num_iteration == 0 else batch_loss / num_iteration
-        return num_iteration, buffer_batch_new, buffer_batch_exp, batch_loss
-
-    def fit(self):
-        print("weak?")
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-
-        self.logger.info("start training...")
-        self._build_buffer_batch()
-        for epo in range(self.start_epoch, epoch_nums): #
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self._scheduler_step()  #
-
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s" \
-                             % (self.epoch_i, loss_total / self.train_batch_nums, train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info(
-                        "---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s" \
-                        % (test_total, test_equ_ac, test_val_ac, test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info(
-                        "---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s" \
-                        % (valid_total, valid_equ_ac, valid_val_ac, valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info(
-                        "---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s" \
-                        % (test_total, test_equ_ac, test_val_ac, test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]''' \
-                         % (self.best_valid_equ_accuracy, self.best_valid_value_accuracy, \
-                            self.best_test_equ_accuracy, self.best_test_value_accuracy))
-
-    def sequence_mask(self, sequence_length, max_len=None):
-        if max_len is None:
-            max_len = sequence_length.data.max()
-        batch_size = sequence_length.size(0)
-        seq_range = torch.arange(0, max_len).long()
-        seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
-        if sequence_length.is_cuda:
-            seq_range_expand = seq_range_expand.cuda()
-        seq_length_expand = (sequence_length.unsqueeze(1).expand_as(seq_range_expand))
-        return seq_range_expand < seq_length_expand
-
-
-class TransformerTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"], config["out_symbol2idx"][SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-        #self.scheduler=torch.optim.lr_scheduler.StepLR(self.optimizer,step_size=5,gamma=0.8)
-        #self.optimizer = WarmUpScheduler(optimizer, self.config["learning_rate"], self.config["embedding_size"], self.config["warmup_steps"])
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _build_loss(self, symbol_size, out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.loss = NLLLoss(weight, pad)
-
-    def _idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.out_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-
-    def _train_batch(self, batch):
-        outputs = self.model(batch["question"], batch["equation"])
-        outputs = torch.nn.functional.log_softmax(outputs, dim=1)
-        if self.config["share_vocab"]:
-            batch_equation = self._idx2word_2idx(batch["equation"])
-            self.loss.eval_batch(outputs, batch_equation.view(-1))
-        else:
-            self.loss.eval_batch(outputs, batch["equation"].view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        test_out = self.model(batch["question"])
-        if self.config["share_vocab"]:
-            target = self._idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            if self.config["task_type"]==TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            elif self.config["task_type"]==TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            else:
-                raise NotImplementedError
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self.optimizer.step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-            # self.logger.info("---------- lr [%1.8f]"%(self.optimizer.get_lr()[0]))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-                            
-
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-
-        test_start_time = time.time()
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-class SeqGANTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"], self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.generator_optimizer=torch.optim.Adam(self.model.generator.parameters(),\
-                                                    lr=self.config["learning_rate"])
-        self.discriminator_optimizer=torch.optim.Adam(self.model.discriminator.parameters(),\
-                                                    lr=self.config["learning_rate"])
-
-    def _build_loss(self, symbol_size, out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.nll_loss = NLLLoss(weight, pad)
-        self.binary_loss = BinaryCrossEntropyLoss()
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "generator_optimizer": self.generator_optimizer.state_dict(),
-            "discriminator_optimizer": self.discriminator_optimizer.state_dict(),
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.generator_optimizer.load_state_dict(check_pnt["generator_optimizer"])
-        self.discriminator_optimizer.load_state_dict(check_pnt["discriminator_optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_value_accuracy = check_pnt["value_acc"]
-        self.best_equ_accuracy = check_pnt["equ_acc"]
-
-    def train_generator(self):
-        print("generator pretrain...")
-        for epo in range(20):
-            loss_total = 0.
-            self.model.generator.train()
-            self.model.discriminator.eval()
-            for batch_idx, batch in enumerate(self.dataloader.load_data("train")):
-                self.model.zero_grad()
-                outputs = self.model.generator.pre_train(batch["question"], batch["ques len"], batch["equation"])
-                #outputs=torch.nn.functional.log_softmax(outputs,dim=1)
-                if self.config["share_vocab"]:
-                    batch_equation = self._idx2word_2idx(batch["equation"])
-                    self.nll_loss.eval_batch(outputs, batch_equation.view(-1))
-                else:
-                    self.nll_loss.eval_batch(outputs, batch["equation"].view(-1))
-                batch_loss = self.nll_loss.get_loss()
-                loss_total += batch_loss
-                self.nll_loss.backward()
-                self.generator_optimizer.step()
-                self.nll_loss.reset()
-            print("epoch [%2d] avr loss [%2.8f]" % (epo + 1, loss_total / self.train_batch_nums))
-
-    def train_discriminator(self):
-        print("discriminator pretrain...")
-        for epo in range(20):
-            loss_total = 0.
-            self.model.generator.eval()
-            self.model.discriminator.train()
-            for batch_idx, batch in enumerate(self.dataloader.load_data("train")):
-                self.model.zero_grad()
-                output, _, _, _ = self.model.generator(batch["question"], batch["ques len"])
-                pred_y = self.model.discriminator(output)
-                label_y = torch.zeros_like(pred_y).to(self.config["device"])
-                self.binary_loss.eval_batch(pred_y, label_y)
-
-                if self.config["share_vocab"]:
-                    batch_equation = self._idx2word_2idx(batch["equation"])
-                else:
-                    batch_equation = batch["equation"]
-                pred_y = self.model.discriminator(batch_equation)
-                label_y = torch.ones_like(pred_y).to(self.config["device"])
-                self.binary_loss.eval_batch(pred_y, label_y)
-
-                norm = self.config['l2_reg_lambda'] * (self.model.discriminator.W_O.weight.norm() + self.model.discriminator.W_O.bias.norm())
-                self.binary_loss.add_norm(norm)
-                loss_total += self.binary_loss.get_loss()
-                self.binary_loss.backward()
-                self.discriminator_optimizer.step()
-                self.binary_loss.reset()
-            print("epoch [%2d] avr loss [%2.8f]" % (epo + 1, loss_total / self.train_batch_nums))
-
-    def get_reward(self, outputs, monte_carlo_outputs, token_logits):
-        rewards = 0
-        batch_size = outputs.size(0)
-        steps = len(monte_carlo_outputs)
-        for idx in range(steps):
-            output = self.model.discriminator(monte_carlo_outputs[idx])
-            reward = output.reshape(batch_size, -1).mean(dim=1)
-            mask = outputs[:, idx] != self.config["out_pad_token"]
-            reward = reward * token_logits[idx] * mask.float()
-            mask_sum = mask.sum()
-            if (mask_sum):
-                rewards += reward.sum() / mask_sum
-        return -rewards
-
-    def _train_batch(self, batch):
-        self.model.generator.train()
-        self.model.discriminator.eval()
-        outputs, _, monte_carlo_outputs, P = self.model.generator(batch["question"], batch["ques len"], batch["equation"])
-        g_loss = self.get_reward(outputs, monte_carlo_outputs, P)
-
-        self.model.generator.eval()
-        self.model.discriminator.train()
-        pred_y = self.model.discriminator(outputs)
-        label_y = torch.zeros_like(pred_y).to(self.config["device"])
-        self.binary_loss.eval_batch(pred_y, label_y)
-
-        if self.config["share_vocab"]:
-            batch_equation = self._idx2word_2idx(batch["equation"])
-        else:
-            batch_equation = batch["equation"]
-        pred_y = self.model.discriminator(batch_equation)
-        label_y = torch.ones_like(pred_y).to(self.config["device"])
-        self.binary_loss.eval_batch(pred_y, label_y)
-
-        norm = self.config['l2_reg_lambda'] * (self.model.discriminator.W_O.weight.norm() + self.model.discriminator.W_O.bias.norm())
-        self.binary_loss.add_norm(norm)
-        d_loss = self.binary_loss.get_loss()
-        self.model.generator.train()
-        self.model.discriminator.train()
-        return g_loss, d_loss
-
-    def _eval_batch(self, batch):
-        test_out = self.model(batch["question"], batch["ques len"])
-        if self.config["share_vocab"]:
-            target = self._idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        g_loss_total = 0.
-        d_loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            g_batch_loss, d_batch_loss = self._train_batch(batch)
-            g_loss_total += g_batch_loss
-            g_batch_loss.backward()
-            self.generator_optimizer.step()
-            d_loss_total += d_batch_loss
-            self.binary_loss.backward()
-            self.discriminator_optimizer.step()
-            self.binary_loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return g_loss_total, d_loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        # generator pretrain
-        self.train_generator()
-        # discriminator pretrain
-        self.train_discriminator()
-        # seqgan train
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            g_loss_total, d_loss_total, train_time_cost = self._train_epoch()
-            print("epoch [%2d] avr g_loss [%2.8f] avr d_loss [%2.8f]"%(self.epoch_i,g_loss_total/self.train_batch_nums,\
-                                                                        d_loss_total/self.train_batch_nums))
-            print("---------- train time {}".format(train_time_cost))
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                equation_ac, value_ac, eval_total, test_time_cost = self.evaluate()
-                print("---------- test equ acc [%2.3f] | test value acc [%2.3f]" % (equation_ac, value_ac))
-                print("---------- test time {}".format(test_time_cost))
-                if value_ac >= self.best_value_accuracy:
-                    self.best_value_accuracy = value_ac
-                    self.best_equ_accuracy = equation_ac
-                    self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-
-class GPT2Trainer(TransformerTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_loss(config["vocab_size"], config["out_symbol2idx"][SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-    
-    def _train_batch(self, batch):
-        outputs, target = self.model(batch["ques_source"], batch["equ_source"])
-        outputs = torch.nn.functional.log_softmax(outputs, dim=1)
-
-        self.loss.eval_batch(outputs, target.view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        test_out, _ = self.model(batch["ques_source"])
-        target = self.model.encode_(batch["equ_source"])
-        batch_size = len(target)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            if self.config["task_type"]==TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            elif self.config["task_type"]==TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-class BERTGenTrainer(TransformerTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_loss(len(self.dataloader.dataset.out_symbol2idx), self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN])
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-
-    # def _build_loss(self, symbol_size, out_pad_token):
-    #     weight = torch.ones(symbol_size).to(self.config["device"])
-    #     pad = out_pad_token
-    #     self.loss = NLLLoss(weight, pad)
-
-    def _train_batch(self, batch):
-        # print (batch)
-        outputs, target = self.model(batch["ques_source"], batch["equ_source"])
-        # print ("outputs, target:", outputs.size(), target.size())
-        outputs = torch.nn.functional.log_softmax(outputs, dim=1)
-
-        # print (outputs.size(), target.size())
-        # print (outputs)
-        # print (target)
-        self.loss.eval_batch(outputs, target.contiguous().view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s" \
-                             % (self.epoch_i, loss_total / self.train_batch_nums, train_time_cost))
-                             # + "\n---------- lr [%1.8f]" % (self.optimizer.get_lr()[0]))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                self.logger.info(
-                    "---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s" \
-                    % (valid_total, valid_equ_ac, valid_val_ac, valid_time_cost))
-                test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                self.logger.info(
-                    "---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s" \
-                    % (test_total, test_equ_ac, test_val_ac, test_time_cost))
-
-                if valid_val_ac >= self.best_valid_value_accuracy:
-                    self.best_valid_value_accuracy = valid_val_ac
-                    self.best_valid_equ_accuracy = valid_equ_ac
-                    self.best_test_value_accuracy = test_val_ac
-                    self.best_test_equ_accuracy = test_equ_ac
-                    self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]''' \
-                         % (self.best_valid_equ_accuracy, self.best_valid_value_accuracy, \
-                            self.best_test_equ_accuracy, self.best_test_value_accuracy))
-
-    def _eval_batch(self, batch):
-        test_out, _ = self.model(batch["ques_source"])
-        target = batch["equ_source"]
-        batch_size = len(target)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            # val_ac, equ_ac, _, _ = self.evaluator.eval_source(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            if self.config["task_type"] == TaskType.SingleEquation:
-
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx],
-                                                             batch["num list"][idx], batch["num stack"][idx], True)
-            elif self.config["task_type"] == TaskType.MultiEquation:
-                # print("test_out[idx]:", test_out[idx])
-                # print("target[idx]:", target[idx])
-                # print()
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], target[idx],
-                                                             batch["num list"][idx], batch["num stack"][idx], True)
-            else:
-                raise NotImplementedError
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-class TRNNTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["temp_symbol_size"], self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN],config["operator_nums"])
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _build_loss(self, symbol_size, out_pad_token, operator_num):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.seq2seq_loss = NLLLoss(weight, pad)
-        weight2=torch.ones(operator_num).to(self.config["device"])
-        self.ans_module_loss=NLLLoss(weight2)
-
-    def _temp_idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.temp_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-    def _equ_idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.out_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-
-    def _idx2symbol(self,equation):
-        symbols=[]
-        eos_idx=self.dataloader.dataset.temp_symbol2idx[SpecialTokens.EOS_TOKEN]
-        pad_idx=self.dataloader.dataset.temp_symbol2idx[SpecialTokens.PAD_TOKEN]
-        sos_idx=self.dataloader.dataset.temp_symbol2idx[SpecialTokens.SOS_TOKEN]
-        for idx in equation:
-            if idx in [eos_idx,pad_idx,sos_idx]:
-                break
-            symbols.append(self.dataloader.dataset.temp_idx2symbol[idx])
-        return symbols
-    
-    def _train_batch_seq2seq(self, batch):
-        outputs = self.model.seq2seq_forward(batch["question"], batch["ques len"], batch["template"])
-        #outputs=torch.nn.functional.log_softmax(outputs,dim=1)
-        if self.config["share_vocab"]:
-            batch_equation = self._temp_idx2word_2idx(batch["template"])
-            self.seq2seq_loss.eval_batch(outputs, batch_equation.view(-1))
-        else:
-            self.seq2seq_loss.eval_batch(outputs, batch["template"].view(-1))
-        batch_loss = self.seq2seq_loss.get_loss()
-        return batch_loss
-    def _train_batch_answer_module(self, batch):
-        '''seq,seq_length,num_pos, template'''
-        for idx,equ in enumerate(batch["equ_source"]):
-            batch["equ_source"][idx]=equ.split(" ")
-        outputs,target = self.model.answer_module_forward(batch["question"], batch["ques len"], batch["num pos"], batch["equ_source"])
-        for b_i in range(len(target)):
-            output=torch.nn.functional.log_softmax(outputs[b_i],dim=1)
-            self.ans_module_loss.eval_batch(output, target[b_i].view(-1))
-        batch_loss = self.ans_module_loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        equ_acc = []
-        val_acc = []
-        template,equation = self.model(batch["question"], batch["ques len"],batch["num pos"])
-        if self.config["share_vocab"]:
-            #temp_target = self._temp_idx2word_2idx(batch["template"])
-            equ_target = self._equ_idx2word_2idx(batch["equation"])
-        else:
-            #temp_target = batch["template"]
-            equ_target = batch["equation"]
-        batch_size = equ_target.size(0)
-        for idx in range(batch_size):
-            # test=self._idx2symbol(template[idx])
-            # tar=self._idx2symbol(temp_target[idx])
-            # if test==tar:
-            #     equ_ac=True
-            # else:
-            #     equ_ac=False
-            if self.config["task_type"]==TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(equation[idx], equ_target[idx], batch["num list"][idx], batch["num stack"][idx])
-            elif self.config["task_type"]==TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(equation[idx], equ_target[idx], batch["num list"][idx], batch["num stack"][idx])
-            else:
-                raise NotImplementedError
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-    def _eval_batch_seq2seq(self,batch):
-        test_out = self.model.seq2seq_forward(batch["question"], batch["ques len"])
-        if self.config["share_vocab"]:
-            target = self._temp_idx2word_2idx(batch["template"])
-        else:
-            target = batch["template"]
-        batch_size = target.size(0)
-        equ_acc = []
-        for idx in range(batch_size):
-            test=self._idx2symbol(test_out[idx])
-            tar=self._idx2symbol(target[idx])
-            if test==tar:
-                equ_ac=True
-            else:
-                equ_ac=False
-            equ_acc.append(equ_ac)
-        return equ_acc
-    def _eval_batch_answer_module(self,batch):
-        for idx,equ in enumerate(batch["equ_source"]):
-            batch["equ_source"][idx]=equ.split(" ")
-        test_out = self.model.test_ans_module(batch["question"], batch["ques len"], batch["num pos"], batch["equ_source"])
-        if self.config["share_vocab"]:
-            target = self._equ_idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        for idx in range(batch_size):
-            if self.config["task_type"]==TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            elif self.config["task_type"]==TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            else:
-                raise NotImplementedError
-            val_acc.append(val_ac)
-        return val_acc
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total_seq2seq = 0.
-        loss_total_ans_module=0.
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            #train seq2seq module
-            self.model.train()
-            self.model.embedder.eval()
-            self.model.attn_encoder.eval()
-            self.model.recursivenn.eval()
-            self.model.zero_grad()
-            batch_loss = self._train_batch_seq2seq(batch)
-            loss_total_seq2seq += batch_loss
-            self.seq2seq_loss.backward()
-            self.optimizer.step()
-            self.seq2seq_loss.reset()
-            #train answer module
-            self.model.train()
-            self.model.seq2seq.eval()
-            self.model.zero_grad()
-            batch_loss = self._train_batch_answer_module(batch)
-            loss_total_ans_module += batch_loss
-            self.ans_module_loss.backward()
-            self.optimizer.step()
-            self.ans_module_loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total_seq2seq,loss_total_ans_module, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total_seq2seq,loss_total_ans_module, train_time_cost = self._train_epoch()
-
-            self.logger.info("epoch [%3d] avr seq2seq module loss [%2.8f] | avr answer module loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total_seq2seq/self.train_batch_nums,loss_total_ans_module/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-    
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(eval_set):
-            #batch_equ_ac = self._eval_batch_seq2seq(batch)
-            #batch_val_ac = self._eval_batch_answer_module(batch)
-            #batch_equ_ac = []
-            batch_val_ac,batch_equ_ac=self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-class Graph2TreeTrainer(GTSTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        # self.logger.info("get group nums...")
-        # self.dataloader.dataset.build_group_nums_for_graph()
-    
-    def _train_batch(self, batch):
-        '''
-        seq, seq_length, nums_stack, num_size, generate_nums, num_pos,\
-                UNK_TOKEN,num_start,target=None, target_length=None,max_length=30,beam_size=5
-        '''
-        unk = self.dataloader.out_unk_token
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in self.dataloader.dataset.generate_list]
-
-        outputs=self.model(batch["question"],batch["ques len"],batch["group nums"],batch["num list"],batch["num stack"],batch["num size"],\
-                                generate_nums,batch["num pos"],num_start,batch["equation"],batch["equ len"],UNK_TOKEN=unk)
-        self.loss.eval_batch(outputs, batch["equation"], batch["equ mask"])
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-    
-    def _eval_batch(self, batch):
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in self.dataloader.dataset.generate_list]
-        test_out=self.model(batch["question"],batch["ques len"],batch["group nums"],batch["num list"],batch["num stack"],batch["num size"],\
-                                generate_nums,batch["num pos"],num_start)
-        
-        if self.config["task_type"]==TaskType.SingleEquation:
-            val_ac, equ_ac, _, _ = self.evaluator.result(test_out, batch["equation"].tolist()[0], batch["num list"][0], batch["num stack"][0])
-        elif self.config["task_type"]==TaskType.MultiEquation:
-            val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out, batch["equation"].tolist()[0], batch["num list"][0], batch["num stack"][0])
-        else:
-            raise NotImplementedError
-        return [val_ac], [equ_ac]
-
-class Graph2TreeIBMTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"],config["out_symbol2idx"][SpecialTokens.PAD_TOKEN])
-        # self.logger.info("build deprel tree...")
-        # self.dataloader.dataset.build_deprel_tree()
-
-    
-    def _build_loss(self, symbol_size,out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.loss = NLLLoss(weight,pad)
-    
-    def _build_optimizer(self):
-        self.encoder_optimizer = torch.optim.Adam(self.model.encoder.parameters(),  lr=self.config["learning_rate"], weight_decay=1e-5)
-        self.decoder_optimizer = torch.optim.Adam(self.model.decoder.parameters(),  lr=self.config["learning_rate"])
-        self.attention_optimizer = torch.optim.Adam(self.model.attention.parameters(),  lr=self.config["learning_rate"])
-    
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "encoder_optimizer": self.encoder_optimizer.state_dict(),
-            "decoder_optimizer": self.decoder_optimizer.state_dict(),
-            "attention_optimizer": self.attention_optimizer.state_dict(), 
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.encoder_optimizer.load_state_dict(check_pnt["encoder_optimizer"])
-        self.decoder_optimizer.load_state_dict(check_pnt["decoder_optimizer"])
-        self.attention_optimizer.load_state_dict(check_pnt["attention_optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-
-    def _optimizer_step(self):
-        self.encoder_optimizer.step()
-        self.decoder_optimizer.step()
-        self.attention_optimizer.step()
-
-    def _model_zero_grad(self):
-        self.model.encoder.zero_grad()
-        self.model.decoder.zero_grad()
-        self.model.attention.zero_grad()
-
-    def _model_train(self):
-        self.model.encoder.train()
-        self.model.decoder.train()
-        self.model.attention.train()
-
-    def _model_eval(self):
-        self.model.encoder.eval()
-        self.model.decoder.eval()
-        self.model.attention.eval()
-
-    def _train_batch(self, batch):
-        outputs,target=self.model(batch["question"],batch["ques len"],batch["group nums"],batch["equation"])
-        self.loss.eval_batch(outputs, target.view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        '''seq, seq_length, group_nums, target'''
-        test_out=self.model(batch["question"],batch["ques len"],batch["group nums"])
-        val_acc = []
-        equ_acc = []
-        batch_size = len(batch["equation"])
-        for idx in range(batch_size):
-            if self.config["task_type"]==TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], batch["equation"][idx], batch["num list"][idx], batch["num stack"][idx])
-            elif self.config["task_type"]==TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[idx], batch["equation"][idx], batch["num list"][idx], batch["num stack"][idx])
-            else:
-                raise NotImplementedError
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-    
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self._model_train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self._model_zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self._optimizer_step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-    
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-        
-
-    def evaluate(self, eval_set):
-        self._model_eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-class MathDQNTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss()
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-    def _build_loss(self):
-        self.loss = MSELoss()
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        #check_pnt = torch.load(self.config["checkpoint_path"],map_location="cpu")
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-
-    def _train_batch(self, batch):
-        '''seq,seq_length,num_pos,target=None'''
-        outputs = self.model(batch['question'],batch["ques len"],batch['num pos'],batch['num list'],batch['ans'],batch['equation'])
-        #outputs = torch.nn.functional.log_softmax(outputs, dim=1)
-        #target = batch['equation']
-        output = outputs[0]
-        target = outputs[1]
-        self.loss.eval_batch(output, target.view(-1))
-        batch_loss = self.loss.get_loss()
-        return batch_loss
-    def _eval_batch(self, batch):
-        '''seq,seq_length,num_pos,target=None'''
-        ans_acc = self.model.predict(batch['question'],batch["ques len"],batch['num pos'],batch['num list'],batch['ans'],batch['equation'])
-        val_acc=[True]*ans_acc+[False]*(len(batch['equation'])-ans_acc)
-        equ_acc=[]
-
-        return val_acc, equ_acc
-    
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.model.train()
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self.optimizer.step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            #self.logger.info(self.model.training)
-            loss_total, train_time_cost = self._train_epoch()
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-            #self.logger.info(self.model.training)
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        #self.logger.info(self.model.training)
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-    
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-            
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-class HMSTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-        self._build_loss(config["symbol_size"],self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN])
-        self.out_pad_idx=self.dataloader.dataset.out_symbol2idx[SpecialTokens.PAD_TOKEN]
-
-    def _build_optimizer(self):
-        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.config["learning_rate"])
-        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=self.config["step_size"], gamma=self.config["scheduler_gamma"])
-    
-    def _build_loss(self,symbol_size,out_pad_token):
-        weight = torch.ones(symbol_size).to(self.config["device"])
-        pad = out_pad_token
-        self.loss = NLLLoss(weight, pad,size_average=False)
-    
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "scheduler": self.scheduler.state_dict(),
-            "start_epoch": self.epoch_i,
-            "best_valid_value_accuracy": self.best_valid_value_accuracy,
-            "best_valid_equ_accuracy": self.best_valid_equ_accuracy,
-            "best_test_value_accuracy": self.best_test_value_accuracy,
-            "best_test_equ_accuracy": self.best_test_equ_accuracy,
-            "best_folds_accuracy": self.best_folds_accuracy,
-            "fold_t":self.config["fold_t"]
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        self.model.load_state_dict(check_pnt["model"])
-        # load parameter of optimizer
-        self.optimizer.load_state_dict(check_pnt["optimizer"])
-        self.scheduler.load_state_dict(check_pnt["scheduler"])
-        # other parameter
-        self.start_epoch = check_pnt["start_epoch"]
-        self.best_valid_value_accuracy = check_pnt["best_valid_value_accuracy"]
-        self.best_valid_equ_accuracy = check_pnt["best_valid_equ_accuracy"]
-        self.best_test_value_accuracy = check_pnt["best_test_value_accuracy"]
-        self.best_test_equ_accuracy = check_pnt["best_test_equ_accuracy"]
-        self.best_folds_accuracy = check_pnt["best_folds_accuracy"]
-    
-    def _idx2word_2idx(self, batch_equation):
-        batch_size, length = batch_equation.size()
-        batch_equation_ = []
-        for b in range(batch_size):
-            equation = []
-            for idx in range(length):
-                equation.append(self.dataloader.dataset.out_symbol2idx[\
-                                            self.dataloader.dataset.in_idx2word[\
-                                                batch_equation[b,idx]]])
-            batch_equation_.append(equation)
-        batch_equation_ = torch.LongTensor(batch_equation_).to(self.config["device"])
-        return batch_equation_
-
-
-    def _train_batch(self, batch):
-        if self.config["share_vocab"]:
-            batch_equation = self._idx2word_2idx(batch["equation"])
-        else:
-            batch_equation = batch["equation"]
-        outputs,_,_ = self.model(batch["spans"], batch["spans len"],batch["span num pos"],batch["word num poses"],batch["span nums"],batch["deprel tree"], batch_equation)
-        #outputs=torch.nn.functional.log_softmax(outputs,dim=1)
-        #outputs = torch.cat(outputs)
-        batch_size = batch_equation.size(0)
-        for step,output in enumerate(outputs):
-            self.loss.eval_batch(output.contiguous().view(batch_size, -1), batch_equation[:, step].view(-1))
-        batch_loss = self.loss.get_loss()
-        total_target_length = (batch_equation != self.out_pad_idx).sum().item()
-        batch_loss = batch_loss / total_target_length
-        return batch_loss
-
-    def _eval_batch(self, batch):
-        test_out = self.model(batch["question"], batch["spans len"],batch["span num pos"],batch["word num poses"],batch["span nums"],batch["deprel tree"])
-        if self.config["share_vocab"]:
-            target = self._idx2word_2idx(batch["equation"])
-        else:
-            target = batch["equation"]
-        batch_size = target.size(0)
-        val_acc = []
-        equ_acc = []
-        for idx in range(batch_size):
-            #val_ac, equ_ac, _, _ = self.evaluator.result(target[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_ac, equ_ac, _, _ = self.evaluator.result(test_out[idx], target[idx], batch["num list"][idx], batch["num stack"][idx])
-            val_acc.append(val_ac)
-            equ_acc.append(equ_ac)
-        return val_acc, equ_acc
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        self.model.train()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            self.batch_idx = batch_idx + 1
-            self.model.zero_grad()
-            batch_loss = self._train_batch(batch)
-            loss_total += batch_loss
-            self.loss.backward()
-            self.optimizer.step()
-            self.loss.reset()
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-        
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s"\
-                                %(self.epoch_i,loss_total/self.train_batch_nums,train_time_cost))
-
-            if epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info("---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s"\
-                                    %(valid_total,valid_equ_ac,valid_val_ac,valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info("---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                    %(test_total,test_equ_ac,test_val_ac,test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]'''\
-                            %(self.best_valid_equ_accuracy,self.best_valid_value_accuracy,\
-                                self.best_test_equ_accuracy,self.best_test_value_accuracy))
-    
-    def evaluate(self, eval_set):
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(eval_set):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-            
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s"\
-                                %(eval_total,equation_ac/eval_total,value_ac/eval_total,test_time_cost))
-
-class SalignedTrainer(AbstractTrainer):
-    def __init__(self, config, model, dataloader, evaluator):
-        super().__init__(config, model, dataloader, evaluator)
-        self._build_optimizer()
-        if config["resume"]:
-            self._load_checkpoint()
-
-        self.min_NUM = self.dataloader.dataset.out_symbol2idx['NUM_0']
-        #print(self.dataloader.dataset.out_symbol2idx); exit()
-        self.do_addeql = False if '<BRG>' in self.dataloader.dataset.out_symbol2idx else True
-        max_NUM = list(self.dataloader.dataset.out_symbol2idx.keys())[-2]
-        self.max_NUM = self.dataloader.dataset.out_symbol2idx[max_NUM]
-        self.ADD = self.dataloader.dataset.out_symbol2idx['+']
-        self.POWER = self.dataloader.dataset.out_symbol2idx['^']
-        self.min_CON = self.N_OPS = self.POWER + 1
-        self.UNK = self.dataloader.dataset.out_symbol2idx['<UNK>']
-        self.max_CON = self.min_NUM - 1
-        self.constant = list(self.dataloader.dataset.out_symbol2idx.keys())[self.min_CON: self.min_NUM]
-        #print('self.constant', self.min_NUM, self.max_NUM, self.ADD, self.POWER, self.min_CON, self.max_CON, self.constant); exit()
-
-        model.encoder.initialize_fix_constant(len(self.constant), self.model._device)
-
-    def _build_optimizer(self):
-        # optimizer
-        self.optimizer = torch.optim.Adam(self.model.parameters(), self.config["learning_rate"],
-                                                   weight_decay=self.config["weight_decay"])
-        # scheduler
-        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=self.config["step_size"],
-                                                               gamma=0.5)
-
-    def _save_checkpoint(self):
-        check_pnt = {
-            "model": self.model.state_dict(),
-        }
-        torch.save(check_pnt, self.config["checkpoint_path"])
-
-    def _load_checkpoint(self):
-        print('load checkpoint ...', self.config["checkpoint_path"]); #exit()
-        check_pnt = torch.load(self.config["checkpoint_path"], map_location=self.config["map_location"])
-        # load parameter of model
-        # print('check_pnt["model"]', check_pnt["model"].keys()) #check_pnt["model"] #
-        pre_trained_dict = check_pnt["model"] #{k: check_pnt["model"][k] for k in check_pnt["model"] if ('decoder.op_selector' not in k and k != '_op_loss.weight')}
-        self.model.load_state_dict(pre_trained_dict, strict=False)
-
-    def _scheduler_step(self):
-        self.scheduler.step()
-
-    def _optimizer_step(self):
-        self.optimizer.step()
-
-    def _model_zero_grad(self):
-        self.model.zero_grad()
-
-    def _model_train(self):
-        self.model.train()
-
-    def _model_eval(self):
-        self.model.eval()
-
-    def _train_batch(self, batch, N_OPS):
-        '''
-        seq, seq_length, nums_stack, num_size, generate_nums, num_pos,\
-                UNK_TOKEN,num_start,target=None, target_length=None,max_length=30,beam_size=5
-        '''
-        unk = self.dataloader.out_unk_token
-        num_start = self.dataloader.dataset.num_start
-        generate_nums = [self.dataloader.dataset.out_symbol2idx[symbol] for symbol in
-                         self.dataloader.dataset.generate_list]
-
-        order = torch.sort(batch['ques len'] * -1)[1]
-        for k in batch:
-            if type(batch[k]) is list:
-                batch[k] = [batch[k][i] for i in order]
-            else:
-                batch[k] = batch[k][order]
-
-        #print('Saligned batch', len(batch), batch.keys())
-        batch_size = len(batch['ques len'])
-        # print(batch['ans'][0], batch['id'], batch['constants'][0], batch['constant_indices'][0])
-        # print(batch['text_len'][0], batch['text'].size(), batch['operations'][0], batch['op_len'][0], batch['indice'][0])
-        #print(batch['ans'][0], batch["id"][0], batch['num list'][0], batch["num pos"][0]);
-        #print(batch["ques len"][0], batch["equation"][0], batch['equ len'][0])
-        #exit()
-        #print('forward')
-        outputs, loss = self.model(batch["question"], batch["equation"], batch["ques len"],
-                                   batch["num pos"], batch['num list'], self.constant, batch["equ len"], N_OPS)
-        #print('train_out', loss, outputs[0], batch["raw_equation"].tolist()[0]); #exit()
-        #print(batch["equ len"], self.evaluator.out_expression_list(batch["raw_equation"].tolist()[0],
-        #                                       batch["num list"][0], copy.deepcopy(batch["num stack"][0])))
-        batch_acc = []
-        for i in range(batch_size):
-            #print('equation', list(batch["equation"][i][1:].cpu().numpy()),  outputs[i]); exit()
-            batch_acc += [int(list(batch["equation"][i][1:].cpu().numpy())[:len(outputs[i])] == outputs[i])]
-        # outputs, loss = self.model(batch["question"], batch["ques len"], batch["num stack"], batch["num size"], \
-        #                      generate_nums, batch["num pos"], num_start, batch["equation"], batch["equ len"],
-        #                      UNK_TOKEN=unk)
-        batch_loss = loss #.get_loss()
-        #print('batch_acc', batch_acc)
-        return batch_loss, batch_acc
-
-    def _eval_batch(self, batch, N_OPS):
-        order = torch.sort(batch['ques len'] * -1)[1]
-        for k in batch:
-            if type(batch[k]) is list:
-                batch[k] = [batch[k][i] for i in order]
-            else:
-                batch[k] = batch[k][order]
-        # if batch["id"]==[80313]:
-        #     print(1)
-        # else:
-        #     return [False], [False]
-        #print(batch["question"])
-        test_out = self.model.predict(batch["question"], batch["equation"], batch["ques len"],
-                              batch["num pos"], batch['num list'], self.constant, batch["equ len"], N_OPS)
-        #print('predict_test_out', test_out[0]["equations_index"]);
-        #print(test_out)
-        #print('test_out', test_out[0]["equations"], batch["equation"].tolist()[0]); #exit()
-        val_acs, equ_acs = [], []
-        for i in range(len(test_out)):
-            if self.config["task_type"] == TaskType.SingleEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result(test_out[i]["equations"], batch["raw_equation"].tolist()[i], batch["num list"][i],
-                                                             batch["num stack"][i])
-            elif self.config["task_type"] == TaskType.MultiEquation:
-                val_ac, equ_ac, _, _ = self.evaluator.result_multi(test_out[i]["equations"], batch["raw_equation"].tolist()[i],
-                                                                   batch["num list"][i], batch["num stack"][i]) #, i==0
-            else:
-                raise NotImplementedError
-            val_acs += [val_ac]
-            equ_acs += [equ_ac]
-        return val_acs, equ_acs
-
-    def adjust_equ(self, op_target, eq_len, num_list):
-        batch_size, batch_len = op_target.size()
-        # change NUM
-        # target_mask = torch.ge(op_target, self.min_NUM) * torch.le(op_target, self.max_NUM).to(torch.long)
-        # op_target = (op_target + self.UNK - self.min_NUM + 4) * target_mask + op_target * (1 - target_mask)
-        # change constants
-        target_mask = torch.ge(op_target, self.min_CON) * torch.le(op_target, self.max_NUM).to(torch.long)
-        op_target = (op_target + 3) * target_mask + op_target * (1 - target_mask)
-        # change unk
-        target_mask = torch.eq(op_target, self.UNK).to(torch.long)
-        op_target = (self.min_NUM + 3) * target_mask + op_target * (1 - target_mask)
-        # change +/-/*//
-        target_mask = torch.ge(op_target, self.ADD) * torch.le(op_target, self.POWER - 1).to(torch.long)
-        op_target = (op_target + 2) * target_mask + op_target * (1 - target_mask)
-        # change padding
-        #print(eq_len, num_list)
-        target_mask = torch.tensor([[1]*eq_len[b]+[0]*(batch_len-eq_len[b]) for b in range(batch_size)]).to(torch.long).to(self.model._device)
-        op_target = op_target * target_mask
-        # attach prefix/postfix
-        batch_size, _ = op_target.size()
-        #if self.do_addeql:
-        eq_postfix = torch.zeros((batch_size, 1), dtype=torch.long).to(self.model._device) + 2
-        op_target = torch.cat([op_target, eq_postfix], dim=1)
-        op_target.scatter_(1, torch.tensor([[idx] for idx in eq_len]).to(self.model._device), self.model.EQL)
-        #op_target[torch.arange(batch_size).unsqueeze(1), eq_len] = self.model.EQL
-        #print('op_target', op_target[:3, :10])
-        gen_var_prefix = [self.min_NUM + len(num) + 3 for num in num_list]
-        #print('gen_var_prefix', self.max_NUM, num_list, gen_var_prefix)
-        gen_var_prefix = torch.tensor(gen_var_prefix, dtype=torch.long).unsqueeze(1).to(self.model._device)
-        #gen_var_prefix = torch.zeros((batch_size, 1), dtype=torch.long).to(self.model._device) + 14 #self.max_NUM + 4
-        x_prefix = torch.zeros((batch_size, 1), dtype=torch.long).to(self.model._device) + self.model.GEN_VAR
-        op_target = torch.cat([x_prefix, gen_var_prefix, op_target], dim=1)
-        #if self.do_addeql:
-        eq_len = [(idx + 3) for idx in eq_len]
-        # else:
-        #     eq_len = [(idx + 2) for idx in eq_len]
-
-        return op_target, eq_len
-
-    def _train_epoch(self):
-        epoch_start_time = time.time()
-        loss_total = 0.
-        acc_total = 0.
-        N_OPS = self.N_OPS
-        #print('self.dataset.out_symbol2idx', self.dataloader.dataset.out_symbol2idx)
-        self._model_train()
-
-        #print('min_NUM, max_NUM', min_NUM, max_NUM, N_OPS); #exit()
-        for batch_idx, batch in enumerate(self.dataloader.load_data(DatasetType.Train)):
-            #print('batch_idx', batch_idx)
-            #if batch_idx != 0: continue
-            #if batch_idx >= 50: continue
-            batch["raw_equation"] = batch["equation"].clone()
-            self.batch_idx = batch_idx + 1
-            self._model_zero_grad()
-            # if self.batch_idx==42:
-            #     continue
-            if True:
-                #print('before adjust', batch["equation"][:3, :30], batch['equ len'])
-                batch["equation"], batch['equ len'] = self.adjust_equ(batch["raw_equation"], batch['equ len'], batch['num list'])
-                #print('after adjust', batch["equation"][:3, :30], batch['equ len']); #exit()
-                batch_loss, batch_acc = self._train_batch(batch, N_OPS)
-                acc_total += np.sum(batch_acc)
-                loss_total += batch_loss
-                batch_loss.backward()
-                self._optimizer_step()
-            else: #except IndexError:
-                pass #print("except IndexError")
-            #self.loss.reset()
-        #exit()
-        #print('acc_total', acc_total)
-        epoch_time_cost = time_since(time.time() - epoch_start_time)
-        return loss_total, epoch_time_cost
-        # print("epoch [%2d]avr loss [%2.8f]"%(self.epoch_i,loss_total /self.batch_nums))
-        # print("epoch train time {}".format(time_since(time.time() -epoch_start_time)))
-
-    def fit(self):
-        train_batch_size = self.config["train_batch_size"]
-        epoch_nums = self.config["epoch_nums"]
-        self.train_batch_nums = int(self.dataloader.trainset_nums / train_batch_size) + 1
-
-        self.logger.info("start training...")
-        for epo in range(self.start_epoch, epoch_nums):
-            self.epoch_i = epo + 1
-            self.model.train()
-            loss_total, train_time_cost = self._train_epoch()
-            self._scheduler_step()
-
-            self.logger.info("epoch [%3d] avr loss [%2.8f] | train time %s" \
-                             % (self.epoch_i, loss_total / self.train_batch_nums, train_time_cost))
-
-            if True: #epo % self.test_step == 0 or epo > epoch_nums - 5:
-                if self.config["k_fold"]:
-                    #print('k_fold'); exit() I changed this !!!
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Train)
-
-                    self.logger.info(
-                        "---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s" \
-                        % (test_total, test_equ_ac, test_val_ac, test_time_cost))
-
-                    if test_val_ac >= self.best_test_value_accuracy:
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-                else:
-                    valid_equ_ac, valid_val_ac, valid_total, valid_time_cost = self.evaluate(DatasetType.Valid)
-
-                    self.logger.info(
-                        "---------- valid total [%d] | valid equ acc [%2.3f] | valid value acc [%2.3f] | valid time %s" \
-                        % (valid_total, valid_equ_ac, valid_val_ac, valid_time_cost))
-                    test_equ_ac, test_val_ac, test_total, test_time_cost = self.evaluate(DatasetType.Test)
-
-                    self.logger.info(
-                        "---------- test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s" \
-                        % (test_total, test_equ_ac, test_val_ac, test_time_cost))
-
-                    if valid_val_ac >= self.best_valid_value_accuracy:
-                        self.best_valid_value_accuracy = valid_val_ac
-                        self.best_valid_equ_accuracy = valid_equ_ac
-                        self.best_test_value_accuracy = test_val_ac
-                        self.best_test_equ_accuracy = test_equ_ac
-                        self._save_model()
-            if epo % 5 == 0:
-                self._save_checkpoint()
-        self.logger.info('''training finished.
-                            best valid result: equation accuracy [%2.3f] | value accuracy [%2.3f]
-                            best test result : equation accuracy [%2.3f] | value accuracy [%2.3f]''' \
-                         % (self.best_valid_equ_accuracy, self.best_valid_value_accuracy, \
-                            self.best_test_equ_accuracy, self.best_test_value_accuracy))
-
-    def evaluate(self, eval_set):
-        self._model_eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-        #print('evaluation set', eval_set)
-        for batch_idx, batch in enumerate(self.dataloader.load_data(eval_set)):
-            #print(batch_idx)
-            if batch_idx >= 100: continue
-            #print('batch', batch.encode('utf8'))
-            batch["raw_equation"] = batch["equation"].clone()
-            batch["equation"], batch['equ len'] = self.adjust_equ(batch["equation"], batch['equ len'], batch['num list'])
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch, self.N_OPS)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-            #print('len(batch_val_ac)', batch_val_ac)
-
-        test_time_cost = time_since(time.time() - test_start_time)
-        return equation_ac / eval_total, value_ac / eval_total, eval_total, test_time_cost
-
-    def test(self):
-        self._load_model()
-        self.model.eval()
-        value_ac = 0
-        equation_ac = 0
-        eval_total = 0
-        test_start_time = time.time()
-
-        for batch in self.dataloader.load_data(DatasetType.Test):
-            batch_val_ac, batch_equ_ac = self._eval_batch(batch, self.N_OPS)
-            value_ac += batch_val_ac.count(True)
-            equation_ac += batch_equ_ac.count(True)
-            eval_total += len(batch_val_ac)
-        test_time_cost = time_since(time.time() - test_start_time)
-        self.logger.info("test total [%d] | test equ acc [%2.3f] | test value acc [%2.3f] | test time %s" \
-                         % (eval_total, equation_ac / eval_total, value_ac / eval_total, test_time_cost))
-
-
-
diff --git a/setup.py b/setup.py
index 238241e..f14a3cc 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,7 @@
 
 setuptools.setup(
     name='mwptoolkit',
-    version='0.0.4',
+    version='0.0.5',
     url="https://github.com/LYH-YF/MWPToolkit",
     author=" ",
     long_description=long_description,