Initial upload of PGPS demo with all dependencies

LM_MODEL.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d0c84cefe6acd4fd66020d40eaebae5da6e3cf231266193d6f53d465ae627d0
+size 64083797

README.md

@@ -1,12 +1,52 @@
 ---
-title: Pgps Demo
-emoji: 📚
-colorFrom: green
-colorTo: pink
+title: PGPS Geometric Problem Solver
+emoji: 📐
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 5.43.1
+sdk_version: 4.16.0
 app_file: app.py
 pinned: false
+license: apache-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PGPS: Neural Geometric Problem Solver Demo
+
+This Space demonstrates the PGPS (Plane Geometry Problem Solver) model, which uses multi-modal neural networks to solve geometry problems.
+
+## How to Use
+
+1. **Upload a Geometry Diagram**: Upload an image containing a geometric diagram (triangles, angles, lines, etc.)
+2. **Enter Problem Text**: Provide the text description of the geometry problem
+3. **Get Solution**: The model will analyze both the diagram and text to generate a solution
+
+## Model Details
+
+- **Architecture**: Multi-modal neural network with visual encoder and text encoder
+- **Task**: Geometric problem solving
+- **Paper**: IJCAI 2023
+- **Original Repository**: [GitHub](https://github.com/mingliangzhang2018/PGPS)
+
+## Features
+
+- Visual diagram parsing
+- Text understanding for geometric problems
+- Expression generation for solutions
+- Support for various geometry problem types
+
+## Limitations
+
+- Best performance with clear, simple geometric diagrams
+- Requires both image and text input for optimal results
+- Limited to plane geometry problems
+
+## Citation
+
+```bibtex
+@inproceedings{zhang2023pgps,
+  title={PGPS: A Neural Geometric Solver},
+  author={Zhang, Mingliang and others},
+  booktitle={IJCAI 2023},
+  year={2023}
+}
+```

app.py

@@ -0,0 +1,244 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+from PIL import Image
+import numpy as np
+import sys
+import os
+
+# Add current directory to path
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from core.network import Network, MLMTransformerPretrain
+from model.backbone import get_visual_backbone
+from model.encoder import get_encoder
+from model.decoder import get_decoder
+from datasets.preprossing import SN
+from datasets.utils import get_combined_text, get_var_arg, get_text_index
+from datasets.operators import normalize_exp
+import datasets.diagram_aug as T_diagram
+
+# Language classes for vocabulary management
+class Lang:
+    def __init__(self):
+        self.word2index = {}
+        self.word2count = {}
+        self.index2word = {0: "PAD", 1: "SOS", 2: "EOS", 3: "UNK"}
+        self.n_words = 4
+        self.class_tag = ['PAD', 'QUE', 'VAR', 'NUM', 'SEP']
+        self.sect_tag = ['PAD', 'TEXT', 'STRU', 'SEM']
+        
+    def add_sentence(self, sentence):
+        for word in sentence.split(' '):
+            self.add_word(word)
+    
+    def add_word(self, word):
+        if word not in self.word2index:
+            self.word2index[word] = self.n_words
+            self.word2count[word] = 1
+            self.index2word[self.n_words] = word
+            self.n_words += 1
+        else:
+            self.word2count[word] += 1
+    
+    def indexes_from_sentence(self, sentence, var_values=None, arg_values=None):
+        indexes = []
+        for word in sentence.split(' '):
+            if word in self.word2index:
+                indexes.append(self.word2index[word])
+            else:
+                indexes.append(3)  # UNK
+        return indexes
+
+# Configuration class
+class Config:
+    def __init__(self):
+        # Visual backbone
+        self.visual_backbone = "ResNet10"
+        self.diagram_size = 128
+        
+        # Encoder
+        self.encoder_type = "gru"
+        self.encoder_layers = 2
+        self.encoder_embedding_size = 256
+        self.encoder_hidden_size = 512
+        self.max_input_len = 400
+        
+        # Decoder
+        self.decoder_type = "rnn_decoder"
+        self.decoder_layers = 2
+        self.decoder_embedding_size = 512
+        self.decoder_hidden_size = 512
+        self.max_output_len = 40
+        
+        # General
+        self.dropout_rate = 0.2
+        self.beam_size = 10
+        self.use_MLM_pretrain = True
+        self.MLM_pretrain_path = './LM_MODEL.pth'
+        self.pretrain_emb_path = ''
+        
+        # Dataset
+        self.without_stru = False
+
+# Initialize model
+def load_model():
+    cfg = Config()
+    
+    # Load vocabularies
+    src_lang = Lang()
+    tgt_lang = Lang()
+    
+    # Load vocab files
+    with open('./vocab/vocab_src.txt', 'r') as f:
+        for line in f:
+            src_lang.add_word(line.strip())
+    
+    with open('./vocab/vocab_tgt.txt', 'r') as f:
+        for line in f:
+            tgt_lang.add_word(line.strip())
+    
+    # Create model
+    model = Network(cfg, src_lang, tgt_lang)
+    
+    # Load pretrained weights if available
+    if os.path.exists('./LM_MODEL.pth'):
+        model.mlm_pretrain.load_model('./LM_MODEL.pth')
+    
+    model.eval()
+    return model, src_lang, tgt_lang, cfg
+
+# Process image and text
+def process_input(image, text_input, model, src_lang, tgt_lang, cfg):
+    # Transform image
+    diagram_transform = T_diagram.Compose([
+        T_diagram.Resize(cfg.diagram_size),
+        T_diagram.CenterCrop(cfg.diagram_size),
+        T_diagram.ToTensor(),
+        T_diagram.Normalize()
+    ])
+    
+    diagram = diagram_transform(image).unsqueeze(0)
+    
+    # Process text input
+    # Create a simple text structure
+    text_sn = SN()
+    text_sn.word_list = text_input.split()
+    text_sn.clause_list = [text_input]
+    
+    # Create empty parsing structures (will be filled with defaults)
+    parsing_stru = SN()
+    parsing_stru.word_list = []
+    parsing_stru.clause_list = []
+    
+    parsing_sem = SN()
+    parsing_sem.word_list = []
+    parsing_sem.clause_list = []
+    
+    # Combine text
+    combine_text = SN()
+    get_combined_text(text_sn, parsing_stru, parsing_sem, combine_text, cfg)
+    
+    # Get text indices
+    text_token, text_sect_tag, text_class_tag = get_text_index(combine_text, src_lang)
+    
+    # Convert to tensors
+    text_dict = {
+        'token': torch.LongTensor([text_token]),
+        'sect_tag': torch.LongTensor([text_sect_tag]),
+        'class_tag': torch.LongTensor([text_class_tag]),
+        'len': torch.LongTensor([len(text_token)])
+    }
+    
+    # Get variables and arguments
+    var_arg_positions, var_values, arg_values = get_var_arg(combine_text, cfg)
+    
+    var_dict = {
+        'pos': torch.LongTensor([var_arg_positions]),
+        'len': torch.LongTensor([len(var_arg_positions)]),
+        'var_value': var_values,
+        'arg_value': arg_values
+    }
+    
+    # Create dummy expression dict for inference
+    exp_dict = {
+        'exp': torch.LongTensor([[1]]),  # SOS token
+        'len': torch.LongTensor([1]),
+        'answer': 0
+    }
+    
+    # Run inference
+    with torch.no_grad():
+        outputs = model(diagram, text_dict, var_dict, exp_dict, is_train=False)
+    
+    # Decode outputs
+    if outputs is not None:
+        # Convert output indices to symbols
+        output_symbols = []
+        for idx in outputs[0]:
+            if idx < len(tgt_lang.index2word):
+                symbol = tgt_lang.index2word[idx]
+                if symbol == 'EOS':
+                    break
+                if symbol not in ['PAD', 'SOS']:
+                    output_symbols.append(symbol)
+        
+        expression = ' '.join(output_symbols)
+        
+        # Try to evaluate the expression
+        try:
+            # Simple evaluation (this would need more sophisticated handling in production)
+            result = eval_expression(expression, var_values, arg_values)
+            return f"Expression: {expression}\nResult: {result}"
+        except:
+            return f"Expression: {expression}\n(Could not evaluate)"
+    
+    return "Could not generate solution"
+
+def eval_expression(expr, var_values, arg_values):
+    # This is a simplified evaluator - would need proper implementation
+    # For now, just return the expression
+    return expr
+
+# Gradio interface
+def predict(image, text):
+    if image is None:
+        return "Please upload a geometry diagram image"
+    
+    if not text:
+        return "Please provide the problem text"
+    
+    try:
+        result = process_input(image, text, model, src_lang, tgt_lang, cfg)
+        return result
+    except Exception as e:
+        return f"Error processing input: {str(e)}"
+
+# Load model on startup
+print("Loading PGPS model...")
+model, src_lang, tgt_lang, cfg = load_model()
+print("Model loaded successfully!")
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="Geometry Diagram"),
+        gr.Textbox(
+            lines=3, 
+            placeholder="Enter the geometry problem text here...\nExample: Find the angle x if angle ABC is 60 degrees",
+            label="Problem Text"
+        )
+    ],
+    outputs=gr.Textbox(label="Solution", lines=5),
+    title="PGPS: Neural Geometric Problem Solver",
+    description="Upload a geometry diagram and provide the problem text to get a solution.",
+    examples=[
+        [None, "Find the value of angle x if angle ABC is 60 degrees and angle BCD is 90 degrees"],
+        [None, "Calculate the area of triangle ABC if AB = 5, BC = 7, and angle B = 60 degrees"]
+    ],
+    theme="default"
+)
+
+if __name__ == "__main__":
+    iface.launch()

config/__init__.py

@@ -0,0 +1,3 @@
+from .config_default import *
+from .logger import *
+

config/config_default.py

@@ -0,0 +1,88 @@
+import argparse
+import torchvision.models as models
+
+
+model_names = sorted(name for name in models.__dict__
+                     if name.islower() and not name.startswith("__") and callable(models.__dict__[name]))
+                   
+criterion_list = ["CrossEntropy", "FocalLoss", "MaskedCrossEntropy"]
+optimizer_list = ["SGD", "ADAM"]
+scheduler_list = ["multistep",'cosine','warmup']
+visual_backbone_list = ['ResNet10', 'mobilenet_v2']
+encoder_list = ['lstm', 'gru', 'transformer']
+decoder_list = ["rnn_decoder", "tree_decoder"]
+eval_method_list = ["completion", "choice", "top3"]
+dataset_list = ['Geometry3K', 'PGPS9K'] 
+
+def get_parser():
+    parser = argparse.ArgumentParser(description='PyTorch PGPS Training')
+    # visual backbone
+    ##############################################################################
+    parser.add_argument('--visual_backbone', default="ResNet10", type=str, choices=visual_backbone_list)
+    parser.add_argument('--diagram_size',  default=128, type=int)
+    # encoder model
+    ##############################################################################
+    parser.add_argument('--encoder_type', default="gru", type=str, choices=encoder_list)
+    parser.add_argument('--encoder_layers', default=2, type=int)
+    parser.add_argument('--encoder_embedding_size', default=256, type=int)
+    parser.add_argument('--encoder_hidden_size', default=512, type=int)
+    parser.add_argument('--max_input_len', default=400, type=int)
+    # decoder model
+    ##############################################################################
+    parser.add_argument('--decoder_type', default="rnn_decoder", type=str, choices=decoder_list)
+    parser.add_argument('--decoder_layers', default=2, type=int)
+    parser.add_argument('--decoder_embedding_size', default=512, type=int)
+    parser.add_argument('--decoder_hidden_size', default=512, type=int)
+    parser.add_argument('--max_output_len', default=40, type=int)
+    # general model
+    ##############################################################################
+    parser.add_argument('--dropout_rate', default=0.2, type=float)
+    parser.add_argument('--beam_size', default=10, type=int)
+    # optimizer
+    ##############################################################################
+    parser.add_argument('--optimizer_type', default="ADAMW", type=str, choices=optimizer_list)
+    parser.add_argument('--lr', default=1e-3, type=float, help='initial learning rate without LM')
+    parser.add_argument('--lr_LM', default=1e-4, type=float, help='initial learning rate of LM')
+    parser.add_argument('--weight_decay', default=0.01, type=float)
+    parser.add_argument('--max_epoch', default=540, type=int)
+    parser.add_argument('--scheduler_type', default="warmup", type=str, choices=scheduler_list)
+    parser.add_argument('--scheduler_step', default=[160, 280, 360, 440, 500], type=list)
+    parser.add_argument('--scheduler_factor', default=0.5, type=float, help='learning rate decay factor')
+    parser.add_argument('--cosine_decay_end', default=0.0, type=float, help='cosine decay end')
+    parser.add_argument('--warm_epoch', default=40, type=int)
+    # criterion
+    ###############################################################################
+    parser.add_argument('--criterion', default="MaskedCrossEntropy", choices=criterion_list, type=str)
+    parser.add_argument('--eval_method', default="top3", choices=eval_method_list, type=str)
+    # dataset      
+    ################################################################################
+    parser.add_argument('--dataset', default="PGPS9K", type=str, choices=dataset_list)
+    parser.add_argument('--dataset_dir', default='./datasets/PGPS9K_all')
+    parser.add_argument('--pretrain_vis_path', default='')
+    parser.add_argument('--vocab_src_path', default='./vocab/vocab_src.txt')
+    parser.add_argument('--vocab_tgt_path', default='./vocab/vocab_tgt.txt')
+    parser.add_argument('--pretrain_emb_path', default='')
+    parser.add_argument('--batch_size', default=128, type=int)
+    parser.add_argument('--random_prob', default=0.5, type=float)
+    parser.add_argument('--without_stru', action='store_true', help='structure clauses are used or not')
+    parser.add_argument('--trim_min_count', default=5, type=int, help='minimum number of word')
+    parser.add_argument('--use_MLM_pretrain', action='store_true', help='use MLM pretrain')
+    parser.add_argument('--MLM_pretrain_path', default='./pretraining_model/LM_MODEL.pth')
+    # print information
+    ###################################################################################
+    parser.add_argument('--dump_path', default="./log/", type=str, help='save log path')
+    parser.add_argument('--print_freq', default=20, type=int, help='print frequency')
+    parser.add_argument('--eval_epoch', default=40, type=int)
+    # general config
+    ###################################################################################
+    parser.add_argument('--workers', default=4, type=int)
+    parser.add_argument('--evaluate_only', action='store_true', help='evaluate model on validation set')
+    parser.add_argument('--resume_model', default="", type=str, help='use pre-trained model')
+    # DistributedDataParallel
+    ###################################################################################
+    parser.add_argument('--local_rank', default=0, type=int, help='node rank for distributed training')
+    parser.add_argument('--init_method', default="env://", type=str, help='distributed init method')
+    parser.add_argument('--debug', action='store_true', help = "if debug than set local rank = 0")
+    parser.add_argument('--seed', default=202302, type=int,help='seed for initializing training. ')
+
+    return parser.parse_args()

config/logger.py

@@ -0,0 +1,27 @@
+import logging
+from logging import handlers
+
+class Logger(object):
+    level_relations = {
+        'debug':logging.DEBUG,
+        'info':logging.INFO,
+        'warning':logging.WARNING,
+        'error':logging.ERROR,
+        'crit':logging.CRITICAL
+    }
+
+    def __init__(self, filename, rank, level='info', when='D', backCount=3, fmt='%(asctime)s - %(levelname)s: %(message)s'):
+        self.logger = logging.getLogger(filename)
+        if rank!=0: return 
+        format_str = logging.Formatter(fmt)
+        self.logger.setLevel(self.level_relations.get(level))
+        sh = logging.StreamHandler()
+        sh.setFormatter(format_str)
+        th = handlers.TimedRotatingFileHandler(filename=filename,when=when,backupCount=backCount,encoding='utf-8')
+        th.setFormatter(format_str)
+        self.logger.addHandler(sh)
+        self.logger.addHandler(th)
+
+def create_logger(filepath, rank):
+    log = Logger(filepath, rank)
+    return log.logger

core/network.py

@@ -0,0 +1,200 @@
+import torch
+import torch.nn as nn
+from model.backbone import get_visual_backbone
+from model.encoder import get_encoder, TransformerEncoder
+from model.decoder import get_decoder
+from utils.utils import *
+import numpy as np
+
+
+class MLMTransformerPretrain(nn.Module):
+
+    def __init__(self, cfg, src_lang):
+        super(MLMTransformerPretrain, self).__init__()
+        self.cfg = cfg
+        self.transformer_en = TransformerEncoder(cfg.encoder_embedding_size)
+        self.text_embedding_src = self.get_text_embedding_src(
+            vocab_size = src_lang.n_words,
+            embedding_dim = cfg.encoder_embedding_size,
+            padding_idx = 0,
+            pretrain_emb_path = cfg.pretrain_emb_path
+        )
+        self.class_tag_embedding = nn.Embedding(
+            len(src_lang.class_tag), 
+            cfg.encoder_embedding_size, 
+            padding_idx=0
+        )
+        self.sect_tag_embedding = nn.Embedding(
+            len(src_lang.sect_tag), 
+            cfg.encoder_embedding_size, 
+            padding_idx=0
+        )
+    
+    def forward(self, text_dict):
+        '''
+            text_dict = {'token', 'sect_tag', 'class_tag', 'len'}
+        '''
+        # text feature
+        token_emb = self.text_embedding_src(text_dict['token'])
+        class_tag_emb = self.class_tag_embedding(text_dict['class_tag'])
+        sect_tag_emb = self.sect_tag_embedding(text_dict['sect_tag'])
+        text_emb_src = token_emb.sum(dim=1) + sect_tag_emb + class_tag_emb
+        transformer_outputs = self.transformer_en(text_dict['len'], text_emb_src)
+        return transformer_outputs
+
+    def load_model(self, model_path):
+        pretrain_dict = torch.load(
+            model_path, map_location="cuda"
+        )
+        pretrain_dict_model = pretrain_dict['state_dict'] \
+                                if 'state_dict' in pretrain_dict else pretrain_dict
+        model_dict = self.state_dict()
+        from collections import OrderedDict
+        new_dict = OrderedDict()
+        for k, v in pretrain_dict_model.items():
+            if k in model_dict:
+                if k.startswith("module"):
+                    new_dict[k[7:]] = v
+                else:
+                    new_dict[k] = v
+        model_dict.update(new_dict)
+        self.load_state_dict(model_dict)
+
+    def get_text_embedding_src(self, vocab_size, embedding_dim, padding_idx, pretrain_emb_path):
+
+        embedding_src = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
+        if pretrain_emb_path!='':
+            emb_content = []
+            with open(pretrain_emb_path, 'r') as f:
+                for line in f:
+                    emb_content.append(line.split()[1:])
+                vector = np.asarray(emb_content, "float32") 
+            embedding_src.weight.data[-len(emb_content):]. \
+                                    copy_(torch.from_numpy(vector))
+        return embedding_src
+
+class Network(nn.Module):
+    
+    def __init__(self, cfg, src_lang, tgt_lang):
+        super(Network, self).__init__()
+        self.cfg = cfg
+        # define the encoder and decoder
+        self.visual_extractor = get_visual_backbone(cfg)  
+        self.encoder = get_encoder(cfg)
+        self.decoder = get_decoder(cfg, tgt_lang)
+        self.visual_emb_unify = nn.ModuleList([
+            nn.Linear(self.visual_extractor.final_feat_dim, cfg.encoder_embedding_size), 
+            nn.ReLU(),
+            nn.Linear(cfg.encoder_embedding_size, cfg.encoder_embedding_size)]
+        )
+        self.visual_emb_unify = nn.Sequential(*self.visual_emb_unify)
+
+        if cfg.use_MLM_pretrain:
+            self.mlm_pretrain = MLMTransformerPretrain(cfg, src_lang)
+            if cfg.MLM_pretrain_path!='':
+                self.mlm_pretrain.load_model(cfg.MLM_pretrain_path)
+        else:
+            self.text_embedding_src = self.get_text_embedding_src(
+                vocab_size = src_lang.n_words,
+                embedding_dim = cfg.encoder_embedding_size,
+                padding_idx = 0,
+                pretrain_emb_path = cfg.pretrain_emb_path
+            )
+            self.class_tag_embedding = nn.Embedding(
+                len(src_lang.class_tag), 
+                cfg.encoder_embedding_size, 
+                padding_idx=0
+            )
+            self.sect_tag_embedding = nn.Embedding(
+                len(src_lang.sect_tag), 
+                cfg.encoder_embedding_size, 
+                padding_idx=0
+            )
+
+        self.src_lang = src_lang
+
+    def forward(self, diagram_src, text_dict, var_dict, exp_dict, is_train=False):
+        '''
+            diagram_src: B x C x W x H
+            text_dict = {'token', 'sect_tag', 'class_tag', 'len'} /
+                        {'token', 'sect_tag', 'class_tag', 'subseq_len', 'item_len', 'item_quant'}
+            var_dict = {'pos', 'len', 'var_value', 'arg_value'}
+            exp_dict = {'exp', 'len', 'answer'}
+        '''
+
+        if self.cfg.use_MLM_pretrain:
+            text_emb_src = self.mlm_pretrain(text_dict)
+        else:
+            # text feature
+            token_emb = self.text_embedding_src(text_dict['token'])
+            class_tag_emb = self.class_tag_embedding(text_dict['class_tag'])
+            sect_tag_emb = self.sect_tag_embedding(text_dict['sect_tag'])
+            # all feature
+            text_emb_src = token_emb.sum(dim=1) + sect_tag_emb + class_tag_emb
+        
+        # diagram feature
+        diagram_emb_src = self.visual_extractor(diagram_src)
+        diagram_emb_src = self.visual_emb_unify(diagram_emb_src).unsqueeze(dim=1)
+        # feature all
+        all_emb_src = torch.cat([diagram_emb_src, text_emb_src], dim=1)
+        text_dict['len'] += 1
+        var_dict['pos'] += 1
+        # encoder
+        encoder_outputs, encode_hidden = self.encoder(all_emb_src, text_dict['len'])
+        problem_output = encode_hidden[-1:,:,:].repeat(self.cfg.decoder_layers, 1, 1)
+        # decoder 
+        outputs = self.decoder(encoder_outputs, problem_output, \
+                                text_dict['len'], \
+                                var_dict['pos'], var_dict['len'], \
+                                exp_dict['exp'], \
+                                is_train)
+        return outputs
+
+    def freeze_module(self, module):
+        self.cfg.logger.info("Freezing module of "+" .......")
+        for p in module.parameters():
+            p.requires_grad = False
+
+    def load_model(self, model_path):
+        pretrain_dict = torch.load(
+            model_path, map_location="cuda"
+        )
+        pretrain_dict_model = pretrain_dict['state_dict'] \
+                            if 'state_dict' in pretrain_dict else pretrain_dict
+        model_dict = self.state_dict()
+        from collections import OrderedDict
+        new_dict = OrderedDict()
+        for k, v in pretrain_dict_model.items():
+            if k.startswith("module"):
+                new_dict[k[7:]] = v
+            else:
+                new_dict[k] = v
+        model_dict.update(new_dict)
+        self.load_state_dict(model_dict)
+        return pretrain_dict
+    
+    def get_text_embedding_src(self, vocab_size, embedding_dim, padding_idx, pretrain_emb_path):
+
+        embedding_src = nn.Embedding(vocab_size, embedding_dim, padding_idx=padding_idx)
+        if pretrain_emb_path!='':
+            emb_content = []
+            with open(pretrain_emb_path, 'r') as f:
+                for line in f:
+                    emb_content.append(line.split()[1:])
+                vector = np.asarray(emb_content, "float32") 
+            embedding_src.weight.data[-len(emb_content):]. \
+                                    copy_(torch.from_numpy(vector))
+        return embedding_src
+
+
+def get_model(args, src_lang, tgt_lang):
+    model = Network(args, src_lang, tgt_lang)
+    args.logger.info(str(model))
+    return model
+
+
+    
+
+
+
+

core/test.py

@@ -0,0 +1,40 @@
+import time
+from utils import *
+
+def validate(args, val_loader, model, tgt_lang):
+    
+    batch_time = AverageMeter('Time', ':5.3f')
+    acc_ans = AverageMeter('Ans_Acc', ':5.4f')
+    acc_eq = AverageMeter('Eq_Acc', ':5.4f')
+    progress = ProgressMeter(len(val_loader), [batch_time, acc_ans, acc_eq], args, prefix='Test: ')
+    # switch to evaluate mode
+    model.eval()
+
+    with torch.no_grad():
+        end = time.time()
+        for i, (diagrams, text_dict, var_dict, exp_dict) in enumerate(val_loader):
+            # set cuda for input data
+            diagrams = diagrams.cuda()
+            set_cuda(text_dict), set_cuda(var_dict), set_cuda(exp_dict)
+            # compute output
+            output = model(diagrams, text_dict, var_dict, exp_dict, is_train=False)
+            if args.eval_method == "completion":
+                acc1, acc2 = compute_exp_result_comp(output, var_dict, exp_dict, tgt_lang)
+            elif args.eval_method == "choice":
+                acc1, acc2 = compute_exp_result_choice(output, var_dict, exp_dict, tgt_lang)
+            elif args.eval_method == "top3":
+                acc1, acc2 = compute_exp_result_topk(output, var_dict, exp_dict, tgt_lang, k_num=3)
+
+            torch.distributed.barrier()
+            
+            reduced_acc_ans = reduce_mean(torch.tensor([acc1]).cuda(), args.nprocs)
+            reduced_acc_eq = reduce_mean(torch.tensor([acc2]).cuda(), args.nprocs)
+
+            acc_ans.update(reduced_acc_ans.item(), len(diagrams))
+            acc_eq.update(reduced_acc_eq.item(), len(diagrams))
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+    return acc_ans.avg, acc_eq.avg

core/train.py

@@ -0,0 +1,44 @@
+import time
+from utils import *
+
+def train(args, epoch, train_loader, model, criterion, optimizer):
+
+    batch_time = AverageMeter('Time', ':5.3f')
+    data_time = AverageMeter('Data', ':5.3f')
+    losses = AverageMeter('Loss', ':.4e')
+    progress = ProgressMeter(len(train_loader), [batch_time, data_time, losses],
+                             args, prefix="Epoch: [{}]".format(epoch))
+
+    # switch to train mode
+    model.train()
+    end = time.time()
+
+    for i, (diagrams, text_dict, var_dict, exp_dict) in enumerate(train_loader):
+        '''
+            text_dict = {'token', 'sect_tag', 'class_tag', 'len'}
+            var_dict = {'pos', 'len', 'var_value', 'arg_value'}
+            exp_dict = {'exp', 'len', 'answer'}
+        '''
+        # measure data loading time
+        data_time.update(time.time() - end)
+        # set cuda for input data
+        diagrams = diagrams.cuda()
+        set_cuda(text_dict), set_cuda(var_dict), set_cuda(exp_dict)
+        # compute output
+        output = model(diagrams, text_dict, var_dict, exp_dict, is_train=True)
+        loss = criterion(output, exp_dict['exp'][:,1:].clone(), exp_dict['len']-1) # Remove special symbol [SOS]
+        # update the loss
+        torch.distributed.barrier()
+        reduced_loss = reduce_mean(loss, args.nprocs)
+        losses.update(reduced_loss.item(), len(diagrams))
+        # compute gradient and do SGD step
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+        if i % args.print_freq == 0:
+            progress.display(i, lr = optimizer.state_dict()['param_groups'][0]['lr'])
+
+    return losses.avg

core/worker.py

@@ -0,0 +1,73 @@
+import torch
+import torch.optim
+import torch.utils.data
+import torch.nn.parallel
+from core.train import *
+from core.test import *
+from utils import *
+from core.network import get_model
+from loss import get_criterion
+from datasets import get_dataloader
+
+
+def main_worker(args):
+
+    args.logger = initialize_logger(args)
+    train_loader, train_sampler, val_loader, src_lang, tgt_lang = get_dataloader(args)
+    model = get_model(args, src_lang, tgt_lang).cuda()
+    optimizer = get_optimizer(args, model)
+    scheduler = get_scheduler(args, optimizer)
+    criterion = get_criterion(args)
+    start_epoch = 0
+    
+    # resume model
+    if not args.resume_model =='':
+        resume_model_dict = model.load_model(args.resume_model)
+        optimizer.load_state_dict(resume_model_dict['optimizer'])
+        scheduler.load_state_dict(resume_model_dict['scheduler'])
+        start_epoch = resume_model_dict["epoch"]+1
+        args.logger.info("The whole model has been loaded from "+ args.resume_model)
+        args.logger.info("The model resumes from epoch "+ str(resume_model_dict["epoch"]))
+        if args.evaluate_only:
+            acc_ans, acc_eq = validate(args, val_loader, model, tgt_lang)
+            args.logger.info("----------Epoch:{:>3d}, test answer_acc {:>5.4f}, equation_acc {:>5.4f} ---------" \
+                                            .format(resume_model_dict["epoch"], acc_ans, acc_eq))
+            return
+    else:
+        args.logger.info("The model is trained from scratch")
+
+    # distributed parallel training 
+    model = torch.nn.parallel.DistributedDataParallel(
+        model, 
+        device_ids=[args.local_rank], 
+        output_device=args.local_rank, 
+        find_unused_parameters=True
+        )
+
+    min_loss = 1e10 
+    
+    for epoch in range(start_epoch, args.max_epoch):
+        # train for one epoch
+        train_sampler.set_epoch(epoch)
+        loss = train(args, epoch, train_loader, model, criterion, optimizer)
+        args.logger.info("----------Epoch:{:>3d}, training loss is {:>5.4f} ---------". \
+                    format(epoch, loss))
+        # evaluate on validation set and save model 
+        if args.local_rank == 0: 
+            if epoch % args.eval_epoch==0 or epoch>=args.max_epoch-5:
+                save_checkpoint({
+                    'epoch': epoch ,
+                    'state_dict': model.state_dict(),
+                    'scheduler': scheduler.state_dict(),
+                    'optimizer': optimizer.state_dict()}, False, args.dump_path)
+            if loss<min_loss: 
+                min_loss = loss
+                save_checkpoint({
+                    'epoch': epoch ,
+                    'state_dict': model.state_dict(),
+                    'scheduler': scheduler.state_dict(),
+                    'optimizer': optimizer.state_dict()}, True, args.dump_path)
+        # learning scheduler step
+        scheduler.step()
+    
+    args.logger.info("------------------- Train Finished -------------------")

datasets/__init__.py

@@ -0,0 +1,37 @@
+from torch.utils.data import DataLoader
+from datasets.dataset import MyDataset
+from torch.utils.data.distributed import DistributedSampler
+from datasets.preprossing import *
+import os
+
+def get_dataloader(args):
+
+    src_lang = SrcLang(args.vocab_src_path)
+    tgt_lang = TgtLang(args.vocab_tgt_path)
+
+    train_data_path = os.path.join(args.dataset_dir, args.dataset, 'train.json')
+    train_pairs = get_raw_pairs(train_data_path)
+    test_data_path = os.path.join(args.dataset_dir, args.dataset, 'test.json')
+    test_pairs = get_raw_pairs(test_data_path)
+
+    train_data = MyDataset(args, train_pairs, src_lang, tgt_lang, is_train=True)
+    train_sampler = DistributedSampler(train_data, shuffle=True)
+    train_loader = DataLoader(dataset=train_data, \
+                              batch_size=int(args.batch_size/args.nprocs), \
+                              pin_memory=True, \
+                              collate_fn=collater(args), \
+                              num_workers=args.workers, \
+                              sampler=train_sampler
+                              )
+                              
+    test_data = MyDataset(args, test_pairs, src_lang, tgt_lang, is_train=False)
+    test_sampler = DistributedSampler(test_data, shuffle=False)
+    test_loader = DataLoader(dataset=test_data, \
+                            batch_size=1, \
+                            pin_memory=True, \
+                            collate_fn=collater(args), \
+                            num_workers=args.workers, \
+                            sampler=test_sampler
+                            )
+                            
+    return train_loader, train_sampler, test_loader, src_lang, tgt_lang

datasets/dataset.py

@@ -0,0 +1,85 @@
+import torch
+import os
+from PIL import Image
+import datasets.diagram_aug as T_diagram
+import datasets.text_aug as T_text
+from datasets.operators import normalize_exp
+from datasets.utils import get_combined_text, get_var_arg, get_text_index
+from datasets.preprossing import SN
+
+class MyDataset(torch.utils.data.Dataset):
+
+    def __init__(self, args, pairs, src_lang, tgt_lang, is_train=True):
+        super().__init__()
+        self.args = args
+        self.pairs = pairs
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+        self.is_train = is_train
+        if is_train: 
+            random_prob = args.random_prob
+        else:
+            random_prob = 0
+        self.diagram_transform = T_diagram.Compose([
+            T_diagram.Resize(args.diagram_size),
+            T_diagram.CenterCrop(args.diagram_size),
+            T_diagram.RandomFlip(random_prob),
+            T_diagram.ToTensor(),
+            T_diagram.Normalize()
+        ])
+        self.text_transform = T_text.Compose([
+            T_text.Point_RandomReplace(random_prob),
+            T_text.AngID_RandomReplace(random_prob),
+            # T_text.Arg_RandomReplace(random_prob),
+            T_text.StruPoint_RandomRotate(random_prob),
+            # T_text.SemPoint_RandomRotate(random_prob),
+            T_text.SemSeq_RandomRotate(random_prob),
+            T_text.StruSeq_RandomRotate(random_prob),
+        ])
+        
+    def __getitem__(self, idx):
+        '''
+            pair{
+                'diagram': str
+                'text': SN()
+                'parsing_stru_seqs': SN()
+                'parsing_sem_seqs': SN()
+                'expression': list
+                'answer': str
+                }
+        '''
+        pair = self.pairs[idx]
+
+        # diagram
+        diagram_path = os.path.join(self.args.dataset_dir, 'Diagram', pair['diagram'])
+        diagram = Image.open(diagram_path).convert("RGB")
+        diagram = self.diagram_transform(diagram)
+        # text, parsing_stru_seqs, parsing_sem_seqs, 
+        self.text_transform(pair['text'], 
+                            pair['parsing_stru_seqs'], 
+                            pair['parsing_sem_seqs'],
+                            pair['expression'])
+        combine_text = SN()
+        get_combined_text(pair['text'], 
+                            pair['parsing_stru_seqs'], 
+                            pair['parsing_sem_seqs'],
+                            combine_text,
+                            self.args)
+        text_token, text_sect_tag, text_class_tag = \
+                            get_text_index(combine_text, self.src_lang)
+        # var and arg
+        var_arg_positions, var_values, arg_values = \
+                            get_var_arg(combine_text, self.args)
+        # expression
+        expression = normalize_exp(pair['expression'])
+        expression = self.tgt_lang.indexes_from_sentence(expression, var_values, arg_values)
+        # choices
+        choices = [float(item) for item in pair['choices']]
+        
+        return  diagram, \
+                text_token, text_sect_tag, text_class_tag, \
+                var_arg_positions, var_values, arg_values, \
+                expression, pair['answer'], pair['id'], choices 
+                
+    def __len__(self):
+        return len(self.pairs)

datasets/diagram_aug.py

@@ -0,0 +1,79 @@
+import random
+from torchvision.transforms import functional as F
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image):
+        for t in self.transforms:
+            image = t(image)
+        return image
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += "    {0}".format(t)
+        format_string += "\n)"
+        return format_string
+
+class Resize(object):
+    '''
+        Resize the training diagram samples, resize the longest edge as max_size
+    '''
+    def __init__(self, max_size):
+        self.max_size = max_size
+
+    def get_size(self, image_size):
+        w, h = image_size
+        if w < h:
+            ow = int(w * self.max_size / h)
+            oh = self.max_size
+        else:
+            ow = self.max_size
+            oh = int(h * self.max_size / w)
+        return (oh, ow)
+
+    def __call__(self, image):
+        size = self.get_size(image.size)
+        image = F.resize(image, size)
+        return image
+
+class CenterCrop(object):
+    '''
+        Crops the given image at the center.
+    '''
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, image):
+        return F.center_crop(image, self.size)
+
+class RandomFlip(object):
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, image):
+        if random.random() < self.prob:
+            flip_method = random.choice([0,1,2])
+            if flip_method==0:
+                image = F.hflip(image)
+            elif flip_method==1:
+                image = F.vflip(image)
+            elif flip_method==2:
+                image = F.vflip(F.hflip(image))
+        return image
+
+class ToTensor(object):
+    def __call__(self, image):
+        return F.to_tensor(image)
+
+class Normalize(object):
+    def __init__(self, mean=[0.85,0.85,0.85], std=[0.3,0.3,0.3]):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image):
+        image = F.normalize(image, mean=self.mean, std=self.std)
+        return image

datasets/operators.py

@@ -0,0 +1,633 @@
+from sympy.parsing.latex import parse_latex
+from sympy.printing import latex
+from sympy import solve
+from sympy.core.numbers import Float
+
+################# Program Executor ########################
+
+spec_token_list = ['frac', 'pi', 'sqrt']
+spec_letter_list = ['f', 'r', 'a', 'c', 'p', 'i', 's', 'q', 'r', 't']
+low_case_list = [chr(i) for i in range(97, 123)]
+fixed_order_ops = [
+    'Get', 'Iso_Tri_Ang', 'Gsin', 'Gcos', 'Gtan', 'Geo_Mean', 'Ratio', 'TanSec_Ang', \
+    'Chord2_Ang', 'Tria_BH_Area', 'Para_Area', 'Kite_Area', 'Circle_R_Circum', \
+    'Circle_D_Circum', 'Circle_R_Area', 'Circle_D_Area', 'ArcSeg_Area', 'Ngon_Angsum', \
+    'RNgon_B_Area', 'RNgon_L_Area', 'RNgon_H_Area']
+alterable_order_ops = [
+    'Sum', 'Multiple', 'Equal', 'Gougu', 'Cos_Law', 'Sin_Law', 'Median', 'Proportion', \
+    'Tria_SAS_Area', 'PRK_Perim', 'Rect_Area', 'Rhom_Area', 'Trap_Area']
+arith_op_list = fixed_order_ops + alterable_order_ops
+priority_list = ["V0", "V1", "V2", "V3", "V4", "V5", "V6", \
+                "N0", "N1", "N2", "N3", "N4", "N5", "N6", "N7", "N8", "N9", "N10", \
+                "C0.5", "C2", "C3", "C4", "C5", "C6", "C8", "C60", "C90", "C180", "C360"]
+V_NUM = 10
+
+class Varible_Record(object):
+    def __init__(self):
+        self.varible_dict = dict()
+        self.mid_varible_dict = dict()
+        self.result = ''
+
+def get_priority(token):
+    if token in priority_list:
+        return priority_list.index(token)
+    else:
+        return -1 # arg
+
+def is_exist_operator(func, ANNO):
+    if not func in arith_op_list:
+        print("Can Not Find Operators!")
+        raise Exception
+    return func
+
+def choose_result(result_list):
+    if len(result_list)==0:
+        return None
+    elif len(result_list)==1:
+        return result_list[0]
+    elif len(result_list)>1:
+        t1 = result_list[0].evalf()
+        t2 = result_list[1].evalf()
+        if (t1>t2 and t2<=0) or (t1<t2 and t1>0):
+            return result_list[0]
+        else:
+            return result_list[1]
+        
+def operand_update(operands, ANNO):
+    for id in range(len(operands)):
+        # Substitute variable
+        if operands[id] in ANNO.mid_varible_dict:
+            operands[id] = "("+ANNO.mid_varible_dict[operands[id]]+")"
+            # pi
+            if "\\pi" in operands[id]:
+                operands[id]=operands[id].replace('\\pi','(3.141593)')
+            # mixed number (improper fraction)
+            if "\\frac" in operands[id]:
+                loc = operands[id].index("\\frac")
+                if loc>0 and operands[id][loc-1].isdigit():
+                    operands[id] = operands[id][:loc]+'+'+operands[id][loc:]
+            continue
+        # Substitute process(intermediate) variable
+        if operands[id] in ANNO.mid_varible_dict:
+            operands[id] = "("+ANNO.mid_varible_dict[operands[id]]+")"
+            continue 
+        # Substitute constant
+        if operands[id][0] == 'C':
+            operands[id] = operands[id][1:]
+
+    return operands
+
+def mid_var_solve(expr_step, ANNO, visit_list, midvar2letter):
+
+    # replace process(intermediate) variable
+    for key, value in midvar2letter.items():
+        expr_step = expr_step.replace(key, value)
+    # Convert the latex form expression to sympy solvable form
+    expr_step = parse_latex(expr_step)
+    # Solving argument
+    for letter in visit_list:
+        try:
+            result = solve(expr_step, letter)
+            result = choose_result(result)
+        except:
+            ANNO.mid_varible_dict[letter] = letter
+            continue
+        if not result is None:
+            result = latex(result)
+            is_update = True
+            result_t = result[:]
+            for item in spec_token_list:
+                result_t = result_t.replace(item, '')
+            # more than one unknown varibles 
+            for letter_t in visit_list:
+                if letter_t in result_t and letter<letter_t:
+                    is_update = False
+                    break
+            # intermediate variables are existed 
+            for key, value in midvar2letter.items():
+                if value in result_t:
+                    is_update = False
+                    break
+            if is_update:
+                ANNO.mid_varible_dict[letter] = result
+            else:
+                ANNO.mid_varible_dict[letter] = letter
+
+    # Solving process(intermediate) variable
+    for key1, value1 in midvar2letter.items():
+        if value1 in str(expr_step):
+            result = solve(expr_step, value1)
+            result = choose_result(result)
+            if not result is None:
+                # Convert the intermediate variable to latex form
+                result = latex(result)
+                # Convert lowercase letters to intermediate variables V_i
+                is_update = True
+                # more than one intermediate variables, only take the front intermediate variables
+                for key2, value2 in midvar2letter.items():
+                    if value2 in result and value1<value2:
+                        is_update = False
+                        break
+                    result.replace(value2, key2)
+                if is_update:
+                    ANNO.mid_varible_dict[key1] = result
+                else:
+                    ANNO.mid_varible_dict[key1] = key1
+
+def mid_var_update(ANNO, visit_list, midvar2letter, midletter2var, is_subs_visit=True):
+
+    has_solved_list = []
+    # Find solved process varibles and arguments
+    for key, value in ANNO.mid_varible_dict.items():
+        if value!='' and isinstance(parse_latex(value).evalf(), Float) or \
+                (key in visit_list and is_subs_visit):
+            if not key in midvar2letter:
+                has_solved_list.append(key)
+            else:
+                has_solved_list.append(midvar2letter[key])
+
+    for key, mid_var in ANNO.mid_varible_dict.items():
+        if value!='' and not key in has_solved_list:
+            # Process varibles V_i are replaced as lowercase letters
+            for key1, value1 in midvar2letter.items():
+                mid_var = mid_var.replace(key1, value1)
+            # Special characters are replaced with '@' for marking
+            mid_var_t = mid_var[:]
+            for item in spec_token_list:
+                mid_var_t = mid_var_t.replace(item, "@"*len(item))
+            # Lowercase letters are replaced with solved values
+            mid_var_new = ''
+            for id in range(len(mid_var_t)):
+                if mid_var_t[id]!="@" and mid_var_t[id] in has_solved_list:
+                    if mid_var_t[id] in midletter2var:
+                        mid_var_new += "("+ANNO.mid_varible_dict[midletter2var[mid_var_t[id]]]+')'
+                    else:
+                        mid_var_new += "("+ANNO.mid_varible_dict[mid_var_t[id]]+')'
+                else:
+                    mid_var_new += mid_var[id]
+            # Lowercase letters are replaced with V_i
+            for key2, value2 in midvar2letter.items():
+                mid_var_new = mid_var_new.replace(value2, key2)
+            ANNO.mid_varible_dict[key] = mid_var_new
+
+def Get(ANNO, arg_list):
+    """
+        Get(a) -> get numerical value of a
+    """
+    if len(arg_list)!=1:
+        print("<Gets> function has only 1 augment!")
+        raise Exception
+
+    if arg_list[0] in ANNO.mid_varible_dict:
+        result = ANNO.mid_varible_dict[arg_list[0]]
+    else:
+        result_v = ANNO.varible_dict[arg_list[0]]
+        result_t = result_v[:]
+        for item in spec_token_list:
+            result_t = result_t.replace(item, "@"*len(item))
+        # Lowercase letters are replaced with solved values
+        result = ''
+        for id in range(len(result_t)):
+            if result_t[id]!="@" and result_t[id] in ANNO.mid_varible_dict:
+                result += "("+ANNO.mid_varible_dict[result_t[id]]+')'
+            else:
+                result += result_v[id]
+    ANNO.result = format(float(parse_latex(result).evalf()),'0.3f')
+
+def Sum(arg_list):
+    """
+        Sum(a, b, c, d) -> a+b+c=d
+    """
+    if len(arg_list)<3:
+        print("<Sum> function has 3 augments at least!")
+        raise Exception
+    expr_step = arg_list[0]
+    for item in arg_list[1:-1]:
+        expr_step += "+" + item
+    expr_step += "-" + arg_list[-1]
+    return expr_step
+
+def Multiple(arg_list):
+    """
+        Multiple(a, b, c, d, e) -> a*b*c*d=e
+    """
+    if len(arg_list)<3:
+        print("<Product> function has 3 augments at least!")
+        raise Exception
+    expr_step = arg_list[0]
+    for item in arg_list[1:-1]:
+        expr_step += "*" + item
+    expr_step += "-" + arg_list[-1]
+    return expr_step
+
+def Equal(arg_list):
+    """
+        Equal(a, b) -> a=b
+    """
+    if len(arg_list)!=2:
+        print("<Equal> function has 2 augments!")
+        raise Exception
+    expr_step = arg_list[0] + "-" + arg_list[-1]
+    return expr_step
+
+def Iso_Tri_Ang(arg_list):
+    """
+        Iso_Tri_Ang(a, b) -> a+2*b=180
+    """
+    if len(arg_list)!=2:
+        print("<Iso_Tri_Ang> function has 2 augments!")
+        raise Exception
+    expr_step = arg_list[0] + "+2*" + arg_list[-1]+"-180"
+    return expr_step
+
+def Gougu(arg_list):
+    """
+        Gougu(a, b, c) -> a^2+b^2=c^2
+    """
+    if len(arg_list)!=3:
+        print("<Gougu> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'^{2}'+"+"+arg_list[1]+"^{2}"+'-'+arg_list[2]+"^{2}"
+    return expr_step
+
+def Gsin(arg_list):
+    """
+        Gsin(a, b, c) -> sin(c)=a/b
+    """
+    if len(arg_list)!=3:
+        print("<Gsin> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'/'+arg_list[1]+'-'+'\\sin{'+arg_list[2]+'/180*3.141593}'
+    return expr_step
+
+def Gcos(arg_list):
+    """
+        Gcos(a, b, c) -> cos(c)=a/b 
+    """
+    if len(arg_list)!=3:
+        print("<Gcos> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'/'+arg_list[1]+'-'+'\\cos{'+arg_list[2]+'/180*3.141593}'
+    return expr_step
+
+def Gtan(arg_list):
+    """
+        Gtan(a, b, c) -> tan(c)=a/b 
+    """
+    if len(arg_list)!=3:
+        print("<Gtan> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'/'+arg_list[1]+'-'+'\\tan{'+arg_list[2]+'/180*3.141593}'
+    return expr_step
+
+def Cos_Law(arg_list):
+    """
+        Cos_Law(a, b, c, d) -> a^2=b^2+c^2-2*b*c
+    """
+    if len(arg_list)!=4:
+        print("<Cos_Law> function has 4 augments!")
+        raise Exception
+    expr_step = arg_list[1]+'^{2}'+"+"+arg_list[2]+"^{2}"+'-'+arg_list[0]+"^{2}"+ \
+                '-'+"2*"+arg_list[1]+'*'+arg_list[2]+'*'+'\\cos{'+arg_list[3]+'/180*3.141593}'
+    return expr_step
+
+def Sin_Law(arg_list):
+    """
+        Sin_Law(a, b, c, d) -> sin(a)/b=sin(c)/d
+    """
+    if len(arg_list)!=4:
+        print("<Sin_Law> function has 4 augments!")
+        raise Exception
+    expr_step = arg_list[3]+'*'+'\\sin{'+arg_list[0]+'/180*3.141593}'+'-'+ \
+                    arg_list[1]+'*'+'\\sin{'+arg_list[2]+'/180*3.141593}'
+    return expr_step
+
+def Median(arg_list):
+    """
+        Median(a, b, c) -> a+c=2*b        
+    """
+    if len(arg_list)!=3:
+        print("<Median> function has 3 augments!")
+        raise Exception 
+    expr_step = arg_list[0]+'-2*'+arg_list[1]+"+"+arg_list[2]
+    return expr_step
+
+def Geo_Mean(arg_list):
+    """
+        Geo_Mean(a, b, c) -> a*b=c^2
+    """
+    if len(arg_list)!=3:
+        print("<Geo_Mean> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[1]+"-"+arg_list[2]+'^{2}'
+    return expr_step
+
+def Proportion(arg_list):
+    """
+        Proportion(a, b, c, d) -> a/b=c/d
+        Proportion(a, b, c, d, e) -> (a/b)^e=c/d
+    """
+    if len(arg_list)<4:
+        print("<Proportion> function has 4 augments at least!")
+        raise Exception
+    if len(arg_list)==4:
+        expr_step = arg_list[0]+'*'+arg_list[3]+"-"+arg_list[1]+'*'+arg_list[2]
+    else:
+        expr_step = arg_list[0]+'*'+arg_list[3]+'^{1/'+arg_list[4]+"}-"+arg_list[1]+'*'+arg_list[2]+'^{1/'+arg_list[4]+"}"
+    return expr_step
+
+def Ratio(arg_list):
+    """
+        Ratio(a, b, c) -> a/b=c
+        Ratio(a, b, c, d) -> (a/b)^c=d
+    """
+    if len(arg_list)<3 or len(arg_list)>4:
+        print("<Power> function has 3 or 4 augments!")
+        raise Exception
+    if len(arg_list)==3:
+        expr_step = arg_list[0]+' / '+arg_list[1]+'-'+arg_list[2]
+    else:
+        expr_step = '('+arg_list[0]+' / '+arg_list[1]+')^{'+arg_list[2]+"}"+"-"+arg_list[3]
+    return expr_step
+
+def Chord2_Ang(arg_list):
+    """
+        Chord2_Ang(a, b, c) -> a=(b+c)/2
+    """
+    if len(arg_list)!=3:
+        print("<Chord2_Ang> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*2-'+arg_list[1]+'-'+arg_list[2]
+    return expr_step
+
+def TanSec_Ang(arg_list):
+    """
+        TanSec_Ang(a, b, c) -> a=(c-b)/2
+    """
+    if len(arg_list)!=3:
+        print("<TanSec_Ang> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*2+'+arg_list[1]+'-'+arg_list[2]
+    return expr_step
+
+def Tria_BH_Area(arg_list):
+    """
+        Tria_BH_Area(a, b, c) -> a*b/2=c
+    """
+    if len(arg_list)!=3:
+        print("<Tria_BH_Area> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[1]+'*0.5-'+arg_list[2]
+    return expr_step
+
+def Tria_SAS_Area(arg_list):
+    """
+        Tria_SAS_Area(a, b, c, d) -> a*c*sin(b)/2=d
+    """
+    if len(arg_list)!=4:
+        print("<Tria_SAS_Area> function has 4 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[2]+'*0.5*\\sin{'+arg_list[1]+'/180*3.141593}-'+arg_list[3]
+    return expr_step
+
+def PRK_Perim(arg_list):
+    """
+        PRK_Perim(a, b, c) -> (a+b)*2=c 
+    """
+    if len(arg_list)!=3:
+        print("<PRK_Perim> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*2+'+arg_list[1]+'*2-'+arg_list[2]
+    return expr_step
+
+def Para_Area(arg_list):
+    """
+        Para_Area(a, b, c) -> a*b=c
+    """
+    if len(arg_list)!=3:
+        print("<Para_Area> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[1]+'-'+arg_list[2]
+    return expr_step
+
+def Rect_Area(arg_list):
+    """
+        Rect_Area(a, b, c) -> a*b=c
+    """
+    if len(arg_list)!=3:
+        print("<Rect_Area> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[1]+'-'+arg_list[2]
+    return expr_step
+
+def Rhom_Area(arg_list):
+    """
+        Rhom_Area(a, b, c) -> a*b*2=c
+    """
+    if len(arg_list)!=3:
+        print("<Phom_Area> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[1]+'*2-'+arg_list[2]
+    return expr_step
+
+def Kite_Area(arg_list):
+    """
+        Kite_Area(a, b, c) -> a*b/2=c
+    """
+    if len(arg_list)!=3:
+        print("<Kite_Area> function has 3 augments!")
+        raise Exception
+    expr_step = arg_list[0]+'*'+arg_list[1]+'*0.5-'+arg_list[2]
+    return expr_step
+
+def Trap_Area(arg_list):
+    """
+        Trap_Area(a, b, c, d) -> (a+b)*c/2=d
+    """
+    if len(arg_list)!=4:
+        print("<Trap_Area> function has 4 augments!")
+        raise Exception
+    expr_step = '0.5*('+arg_list[0]+'+'+arg_list[1]+')*'+arg_list[2]+'-'+arg_list[3]
+    return expr_step
+
+def Circle_R_Circum(arg_list):
+    """
+        Circle_R_Circum(a, b) -> 2*pi*a=b
+        Circle_R_Circum(a, b, c) -> 2*pi*a*b/360=c
+    """
+    if len(arg_list)<2 or len(arg_list)>3:
+        print("<Circle_Circum> function has 2 or 3 augments!")
+        raise Exception
+    if len(arg_list)==2:
+        expr_step = '2*3.141593*'+arg_list[0]+'-'+arg_list[1]
+    else:
+        expr_step = '2*3.141593*'+arg_list[0]+'*'+arg_list[1]+'/360'+'-'+arg_list[2]
+    return expr_step
+
+def Circle_D_Circum(arg_list):
+    """
+        Circle_D_Circum(a, b) -> pi*a=b
+        Circle_D_Circum(a, b, c) -> pi*a*b/360=c
+    """
+    if len(arg_list)<2 or len(arg_list)>3:
+        print("<Circle_Circum> function has 2 or 3 augments!")
+        raise Exception
+    if len(arg_list)==2:
+        expr_step = '3.141593*'+arg_list[0]+'-'+arg_list[1]
+    else:
+        expr_step = '3.141593*'+arg_list[0]+'*'+arg_list[1]+'/360'+'-'+arg_list[2]
+    return expr_step
+
+def Circle_R_Area(arg_list):
+    """
+        Circle_R_Area(a, b) -> pi*a^2=b
+        Circle_R_Area(a, b, c) -> pi*a^2*b/360=c
+    """
+    if len(arg_list)<2 and len(arg_list)>3:
+        print("<Circle_Area> function has 2 or 3 augments!")
+        raise Exception
+    if len(arg_list)==2:
+        expr_step = '3.141593*'+arg_list[0]+'^{2}-'+arg_list[1]
+    else:
+        expr_step = '3.141593*'+arg_list[0]+'^{2}*'+arg_list[1]+'/360'+'-'+arg_list[2]
+    return expr_step
+
+def Circle_D_Area(arg_list):
+    """
+        Circle_D_Area(a, b) -> pi*(a/2)^2=b
+        Circle_D_Area(a, b, c) -> pi*(a/2)^2*b/360=c
+    """
+    if len(arg_list)<2 and len(arg_list)>3:
+        print("<Circle_Area> function has 2 or 3 augments!")
+        raise Exception
+    if len(arg_list)==2:
+        expr_step = '0.25*3.141593*'+arg_list[0]+'^{2}-'+arg_list[1]
+    else:
+        expr_step = '0.25*3.141593*'+arg_list[0]+'^{2}*'+arg_list[1]+'/360'+'-'+arg_list[2]
+    return expr_step
+
+def ArcSeg_Area(arg_list):
+    """
+        ArcSeg_Area(a, b, c) -> pi*a^2*b/360 - a^2*sin(b)/2 = c
+    """
+    if len(arg_list)!=3:
+        print("<ArcSeg_Area> function has 3 augments!")
+        raise Exception
+    expr_step = '3.141593*'+arg_list[0]+'^{2}*'+arg_list[1]+'/360-0.5*'+ \
+                    arg_list[0]+'^{2}*\\sin{'+arg_list[1]+'/180*3.141593}-'+arg_list[2]
+    return expr_step
+
+def Ngon_Angsum(arg_list):
+    """
+        Ngon_Ang(a, b) -> (a-2)*180=b 
+    """
+    if len(arg_list)!=2:
+        print("<Ngon_Ang> function has 2 augments!")
+        raise Exception 
+    expr_step = '('+arg_list[0]+'-2)*180-'+arg_list[1]
+    return expr_step
+
+def RNgon_B_Area(arg_list):
+    """
+        RNgon_B_Area(a, b, c) -> a*b^2/tan(180/a)/4=c
+    """
+    if len(arg_list)!=3:
+        print("<RNgon_B_Area> function has 3 augments!")
+        raise Exception 
+    expr_step = arg_list[0]+'*'+arg_list[1]+'^{2}/4/\\tan{3.141593/'+arg_list[0]+'}-'+arg_list[2]
+    return expr_step
+
+def RNgon_L_Area(arg_list):
+    """
+        RNgon_L_Area(a, b, c) -> a*b^2*sin(360/a)/2=c
+    """
+    if len(arg_list)!=3:
+        print("<RNgon_L_Area> function has 3 augments!")
+        raise Exception 
+    expr_step = arg_list[0]+'*'+arg_list[1]+'^{2}*0.5*\\sin{2*3.141593/'+arg_list[0]+'}-'+arg_list[2]
+    return expr_step
+
+def RNgon_H_Area(arg_list):
+    """
+        RNgon_H_Area(a, b, c) -> a*b^2*tan(180/a)=c
+    """
+    if len(arg_list)!=3:
+        print("<RNgon_H_Area> function has 3 augments!")
+        raise Exception 
+    expr_step = arg_list[0]+'*'+arg_list[1]+'^{2}*\\tan{3.141593/'+arg_list[0]+'}-'+arg_list[2]
+    return expr_step
+
+def result_compute(num_all_list, exp_tokens):
+    ANNO = Varible_Record()
+    # Obtain the mapping between lowercase letters to intermediate variables V_i
+    visit_list = [] # arguments denoted by lowercase letters
+    for num in num_all_list:
+        for item in spec_token_list:
+            num = num.replace(item, "@"*len(item))
+        for letter in num:
+            if letter in low_case_list: visit_list.append(letter)
+    for id, var in enumerate(num_all_list):
+        ANNO.varible_dict["N"+str(id)] = var
+        ANNO.mid_varible_dict["N"+str(id)] = var
+    visit_list.sort() 
+    no_visit_list = list(set(low_case_list)-set(spec_letter_list)-set(visit_list))
+    no_visit_list.sort() # lowercase letters which have not used
+    # mapping between letters to intermediate variables V_i
+    midvar2letter = dict() 
+    midletter2var = dict() 
+    for id in range(V_NUM):
+        midvar2letter['V'+str(id)] = no_visit_list[id]
+        midletter2var[no_visit_list[id]] = 'V'+str(id)
+    # step split
+    step_list = []
+    last_op_id = 0
+    for id, token in enumerate(exp_tokens):
+        if token in arith_op_list and id>0:
+            step_list.append(exp_tokens[last_op_id:id])
+            last_op_id = id
+    step_list.append(exp_tokens[last_op_id:])
+    # run step 
+    for id, step in enumerate(step_list):
+        operator = is_exist_operator(step[0], ANNO)
+        if operator!='Get':
+            operands = operand_update(step[1:], ANNO)
+            expr_step = eval(operator)(operands)
+            mid_var_solve(expr_step, ANNO, visit_list, midvar2letter)
+            mid_var_update(ANNO, visit_list, midvar2letter, midletter2var, True)
+            mid_var_update(ANNO, visit_list, midvar2letter, midletter2var, False)
+        else: 
+            Get(ANNO, step[1:])
+    
+    return ANNO.result
+
+def normalize_exp(exp):
+    # step split
+    step_list = []
+    last_op_id = 0
+    for id, token in enumerate(exp):
+        if token in arith_op_list and id>0:
+            step_list.append(exp[last_op_id:id])
+            last_op_id = id
+    step_list.append(exp[last_op_id:])
+    # normalize step
+    new_exp = []
+    for step in step_list:
+        if step[0] in alterable_order_ops:
+            if step[0] in ['Sum', 'Multiple']: 
+                begin_id, end_id = 1, -1
+                step[begin_id: end_id] = sorted(step[begin_id: end_id], key=lambda token:get_priority(token))
+            if step[0] in ['Equal', 'Gougu', 'PRK_Perim', 'Rect_Area', 'Rhom_Area', 'Trap_Area']: 
+                begin_id, end_id = 1, 3
+                step[begin_id: end_id] = sorted(step[begin_id: end_id], key=lambda token:get_priority(token))
+            if step[0] == 'Cos_Law': 
+                begin_id, end_id = 2, 4
+                step[begin_id: end_id] = sorted(step[begin_id: end_id], key=lambda token:get_priority(token))
+            if step[0] in ['Sin_Law', 'Proportion']:
+                if get_priority(step[1])>get_priority(step[3]) and len(step)==5:
+                    step[1:3], step[3:5] = step[3:5], step[1:3]
+            if step[0] in ['Tria_SAS_Area', 'Median']: 
+                if get_priority(step[1])>get_priority(step[3]):
+                    step[1], step[3] = step[3], step[1]
+        new_exp += step
+
+    return new_exp

datasets/preprossing.py

@@ -0,0 +1,201 @@
+import torch
+import json
+from datasets.utils import *
+
+class SrcLang:
+
+    def __init__(self, vocab_path):
+        self.word2index = {}
+        self.word2count = {}
+        self.index2word = []
+        self.n_words = 0
+        self.get_vocab(vocab_path)
+        self.class_tag = ['[PAD]', '[GEN]', '[POINT]', '[NUM]', '[ARG]', '[ANGID]']
+        self.sect_tag = ['[PAD]', '[PROB]', '[COND]', '[STRU]']
+        
+    def get_vocab(self, vocab_path):
+        with open(vocab_path, 'r') as f:
+            for id, line in enumerate(f):
+                vocab_token = line[:-1]
+                self.word2index[vocab_token] = id
+                self.word2count[vocab_token] = 0
+                self.index2word.append(vocab_token)
+        self.n_words = len(self.index2word)
+    
+    def indexes_from_sentence(self, sentence, id_type='text'):
+        res = []
+        if id_type == 'text':
+            for word in sentence:
+                if word in self.word2index:
+                    res.append(self.word2index[word])
+                    self.word2count[word] += 1
+                else:
+                    res.append(self.word2index["[UNK]"])
+                    self.word2count["[UNK]"] += 1
+                    print("Can not find", word, 'in the src vocab')
+        elif id_type=='class_tag':
+            for word in sentence: res.append(self.class_tag.index(word))
+        elif id_type=='sect_tag':
+            for word in sentence: res.append(self.sect_tag.index(word))
+        return res
+    
+    def sentence_from_indexes(self, indexes):
+        res = []
+        for index in indexes:
+            if index<len(self.index2word):
+                res.append(self.index2word[index])
+            else:
+                res.append("")
+        return res
+
+class TgtLang:
+
+    def __init__(self, vocab_path):
+        self.word2index = {}
+        self.word2count = {}
+        self.index2word = []
+        self.n_words = 0 
+        self.var_start = 0
+        self.get_vocab(vocab_path)
+    
+    def get_vocab(self, vocab_path):
+        spe_num = midvar_num = const_num = 0
+        op_num = var_num = 0
+        
+        with open(vocab_path, 'r') as f:
+            for id, line in enumerate(f):
+                vocab_token = line[:-1]
+                self.word2index[vocab_token] = id
+                self.word2count[vocab_token] = 0
+                self.index2word.append(vocab_token)
+                if vocab_token[0]=='[' and vocab_token[-1]==']': 
+                    spe_num += 1
+                elif vocab_token[0]=='V' and vocab_token[1].isdigit(): 
+                    midvar_num += 1
+                elif vocab_token[0]=='C' and vocab_token[1].isdigit(): 
+                    const_num += 1
+                elif vocab_token[0]=='N' and vocab_token[1].isdigit():
+                    var_num += 1
+                else:
+                    op_num += 1
+
+        self.n_words = len(self.index2word)
+        self.var_start = spe_num + midvar_num + const_num + op_num
+
+    def indexes_from_sentence(self, sentence, var_values, arg_values):
+        res = []
+        for word in sentence:
+            if word in self.word2index:
+                res.append(self.word2index[word])
+                self.word2count[word] += 1
+            elif len(word)==1 and word.islower(): # arg
+                res.append(self.var_start+len(var_values)+arg_values.index(word))
+            else:
+                print("Can not find", word, 'in the tgt vocab')
+        res = [self.word2index["[SOS]"]]+res+[self.word2index["[EOS]"]]
+        return res
+    
+    def sentence_from_indexes(self, indexes, change_dict={}):
+        res = []
+        for index in indexes:
+            if index<len(self.index2word):
+                item = self.index2word[index]
+            else:
+                item = ''
+            if item in change_dict: item = change_dict[item] # var2arg
+            res.append(item)
+        return res
+    
+class SN:
+    def __init__(self):
+        self.token = [] # str list
+        self.sect_tag = [] # [PROB]/[COND]/[STRU]
+        self.class_tag = [] # [GEN]/[NUM]/[ARG]/[POINT]/[ANGID]
+
+def get_raw_pairs(dataset_path):
+
+    raw_pairs = []
+
+    with open(dataset_path, 'r')as fp:
+        content_all = json.load(fp)
+
+    for key, content in content_all.items():
+        text = content['text']
+        stru_seqs = content['parsing_stru_seqs']
+        sem_seqs = content['parsing_sem_seqs']
+        text_data, stru_data, sem_data = SN(), SN(), SN()
+        # tokenization
+        text_data.token = get_token(text)
+        stru_data.token = [get_token(item)+[','] for item in stru_seqs]
+        sem_data.token = [get_token(item)+[','] for item in sem_seqs]
+        # split prob and cond
+        text_data.sect_tag = []
+        stru_data.sect_tag = [['[STRU]']*len(item) for item in stru_data.token]
+        sem_data.sect_tag = [['[COND]']*len(item) for item in sem_data.token]
+        split_text(text_data)
+        # get class tag
+        text_data.class_tag = ['[GEN]']*len(text_data.token)
+        stru_data.class_tag = [['[GEN]']*len(item) for item in stru_data.token]
+        sem_data.class_tag = [['[GEN]']*len(item) for item in sem_data.token]
+        get_point_angleID_tag(text_data, stru_data, sem_data)
+        get_num_arg_tag(text_data, sem_data)
+        # Tag the repeat [NUM] in sem_data which has exist in text_data
+        expression = content['expression'].split(' ')
+        remove_sem_dup(text_data, sem_data, expression)
+
+        content['text'] = text_data
+        content['parsing_stru_seqs'] = stru_data
+        content['parsing_sem_seqs'] = sem_data
+        content['expression'] = expression
+        content['id'] = key
+        
+        raw_pairs.append(content)
+        
+    return raw_pairs
+
+class collater():
+
+    def __init__(self, args):
+        self.args = args
+
+    def __call__(self, batch_data, padding_id=0):
+        diagrams, \
+        text_tokens, text_sect_tags, text_class_tags, \
+        var_arg_positions, var_values, arg_values, \
+        expression, answer, pair_ids, choices  = list(zip(*batch_data))
+        #######################################
+        diagrams = torch.stack(diagrams, dim=0)
+        #######################################
+        len_exp = [len(seq_exp) for seq_exp in expression]
+        max_len_exp = max(len_exp)
+        expression = [seq_exp+[padding_id]*(max_len_exp-len(seq_exp)) for seq_exp in expression]
+        exp_dict = {'exp': torch.LongTensor(expression), 
+                    'len': torch.LongTensor(len_exp),
+                    'answer': answer,
+                    'id': pair_ids,
+                    'choices': choices
+                    }
+        #######################################
+        len_var = [max(len(seq_var),1) for seq_var in var_arg_positions]
+        max_len_var = max(len_var)
+        var_arg_positions = [seq_var+[padding_id]*(max_len_var-len(seq_var)) for seq_var in var_arg_positions]
+        var_dict = {'pos':torch.LongTensor(var_arg_positions),
+                    'len': torch.LongTensor(len_var),
+                    'var_value': var_values,
+                    'arg_value': arg_values
+                    }
+        ########################################
+        len_text = [len(seq_tag) for seq_tag in text_class_tags]
+        max_len_text = max(len_text)
+        for k in range(len(text_tokens)):
+            for j in range(len(text_tokens[k])):
+                text_tokens[k][j] += [padding_id]*(max_len_text-len(text_tokens[k][j]))
+        text_sect_tags = [seq_tag+[padding_id]*(max_len_text-len(seq_tag)) for seq_tag in text_sect_tags]
+        text_class_tags = [seq_tag+[padding_id]*(max_len_text-len(seq_tag)) for seq_tag in text_class_tags]
+        text_dict = {'token': torch.LongTensor(text_tokens),
+                    'sect_tag': torch.LongTensor(text_sect_tags),
+                    'class_tag': torch.LongTensor(text_class_tags),
+                    'len': torch.LongTensor(len_text)
+                    }
+
+        return diagrams, text_dict, var_dict, exp_dict

datasets/text_aug.py

@@ -0,0 +1,233 @@
+import random
+
+upper_case_list = [chr(i) for i in range(65, 91)]
+low_case_list = [chr(i) for i in range(97, 123)]
+angle_id_list = [str(i) for i in range(1, 21)]
+spec_token_list = ['frac', 'pi', 'sqrt']
+
+class Compose(object):
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        for t in self.transforms:
+            t(text_seq, stru_seqs, sem_seqs, exp)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += "    {0}".format(t)
+        format_string += "\n)"
+        return format_string
+
+class Point_RandomReplace(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def get_point_map(self):
+        value_list = [chr(i) for i in range(65, 91)]
+        random.shuffle(value_list)
+        map_dict = {key:value for key, value in zip(upper_case_list, value_list)} 
+        return map_dict        
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        if random.random() < self.prob:
+            map_dict = self.get_point_map()
+            for k in range(len(text_seq.token)):
+                if text_seq.class_tag[k] == '[POINT]':
+                    text_seq.token[k] = map_dict[text_seq.token[k][0]]
+            for k in range(len(stru_seqs.token)):
+                for j in range(len(stru_seqs.token[k])):
+                    if stru_seqs.class_tag[k][j] == '[POINT]':
+                        stru_seqs.token[k][j] = map_dict[stru_seqs.token[k][j][0]]
+            for k in range(len(sem_seqs.token)):
+                for j in range(len(sem_seqs.token[k])):
+                    if sem_seqs.class_tag[k][j] == '[POINT]':
+                        sem_seqs.token[k][j] = map_dict[sem_seqs.token[k][j][0]]
+
+class AngID_RandomReplace(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def get_angid_map(self):
+        value_list = [str(i) for i in range(1, 21)]
+        random.shuffle(value_list)
+        map_dict = {key:value for key, value in zip(angle_id_list, value_list)} 
+        return map_dict        
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        if random.random() < self.prob:
+            map_dict = self.get_angid_map()
+            for k in range(len(text_seq.token)):
+                if text_seq.class_tag[k] == '[ANGID]':
+                    text_seq.token[k] = map_dict[text_seq.token[k]]
+            for k in range(len(sem_seqs.token)):
+                for j in range(len(sem_seqs.token[k])):
+                    if sem_seqs.class_tag[k][j] == '[ANGID]':
+                        sem_seqs.token[k][j] = map_dict[sem_seqs.token[k][j]]
+
+class Arg_RandomReplace(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def get_arg_map(self):
+        value_list = [chr(i) for i in range(97, 123)]
+        random.shuffle(value_list)
+        map_dict = {key:value for key, value in zip(low_case_list, value_list)} 
+        return map_dict
+    
+    def map_arg_in_num(self, map_dict, num):
+        num_t = num[:]
+        new_num = ''
+        for item in spec_token_list:
+            num_t = num_t.replace(item, "@"*len(item))
+        for k in range(len(num_t)):
+            if num_t[k]!='@' and num[k] in low_case_list:
+                new_num += map_dict[num[k]]
+            else:
+                new_num += num[k]
+        return new_num
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        if random.random() < self.prob:
+            map_dict = self.get_arg_map()
+            for k in range(len(text_seq.token)): 
+                if text_seq.class_tag[k] == '[NUM]':
+                    text_seq.token[k] = self.map_arg_in_num(map_dict, text_seq.token[k])
+                if text_seq.class_tag[k] == '[ARG]':
+                    text_seq.token[k] = map_dict[text_seq.token[k]]
+            for k in range(len(sem_seqs.token)):
+                for j in range(len(sem_seqs.token[k])):
+                    if sem_seqs.class_tag[k][j] == '[NUM]':
+                        sem_seqs.token[k][j] = self.map_arg_in_num(map_dict, sem_seqs.token[k][j])
+            for k in range(len(exp)):
+                if exp[k] in low_case_list:
+                    exp[k] = map_dict[exp[k]]
+
+class StruPoint_RandomRotate(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def get_seq_points(self, class_tag):
+        id_list = []
+        begin_point_id = end_point_id = None
+        for id, token in enumerate(class_tag):
+            if token == '[POINT]':
+                if begin_point_id is None:
+                    begin_point_id = id
+            elif not begin_point_id is None and end_point_id is None:
+                end_point_id = id 
+                id_list.append([begin_point_id, end_point_id])
+                begin_point_id = end_point_id = None
+        if not begin_point_id is None and end_point_id is None:
+            id_list.append([begin_point_id, len(class_tag)])
+        
+        return id_list[-1][0], id_list[-1][1]
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        for k in range(len(stru_seqs.token)):
+            if random.random() < self.prob:
+                begin_id, end_id = self.get_seq_points(stru_seqs.class_tag[k])
+                # point on line
+                if stru_seqs.token[k][0] == 'line':
+                    stru_seqs.token[k][begin_id:end_id] = stru_seqs.token[k][end_id-1:begin_id-1:-1]
+                # point on circle
+                if stru_seqs.token[k][0] == '\\odot':
+                    # clockwise change 
+                    if random.random() < 0.5: 
+                        stru_seqs.token[k][begin_id:end_id] = stru_seqs.token[k][end_id-1:begin_id-1:-1]
+                    # set initial point
+                    init_loc = random.randint(begin_id, end_id-1) 
+                    stru_seqs.token[k][begin_id:end_id] = stru_seqs.token[k][init_loc:end_id] + \
+                                                                stru_seqs.token[k][begin_id:init_loc]
+
+class SemPoint_RandomRotate(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def get_seq_points(self, class_tag):
+        id_list = []
+        begin_point_id = end_point_id = None
+        for id, token in enumerate(class_tag):
+            if token == '[POINT]':
+                if begin_point_id is None:
+                    begin_point_id = id
+            elif not begin_point_id is None and end_point_id is None:
+                end_point_id = id 
+                id_list.append((begin_point_id, end_point_id-1))
+                begin_point_id = end_point_id = None
+        if not begin_point_id is None and end_point_id is None:
+            id_list.append((begin_point_id, len(class_tag)-1))
+        
+        return id_list
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        if random.random() < self.prob:
+            for k in range(len(sem_seqs.token)):
+                id_list = self.get_seq_points(sem_seqs.class_tag[k])
+                for begin_id, end_id in id_list:
+                    if random.random() < self.prob:  
+                        sem_seqs.token[k][begin_id], sem_seqs.token[k][end_id] = \
+                            sem_seqs.token[k][end_id], sem_seqs.token[k][begin_id]
+
+class SemSeq_RandomRotate(object):
+
+    def __init__(self, prob=0.5):
+        if prob==0:
+            self.prob = 0 
+        else:    
+            self.prob = prob + 0.2
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        if random.random() < self.prob:
+            # varible id 
+            num_all_list, num_sem_list, num_map_list = [], [], []
+            for item in text_seq.class_tag:
+                if item=='[NUM]':
+                    var_name = 'N'+str(len(num_all_list))
+                    num_all_list.append(var_name)
+                    num_map_list.append(var_name)
+            for k in range(len(sem_seqs.token)):
+                if sem_seqs.class_tag[k][-2] == '[NUM]':
+                    var_name = 'N'+str(len(num_all_list))
+                    num_all_list.append(var_name)
+                    num_sem_list.append([var_name])
+                else:
+                    num_sem_list.append([])
+            # shuffle sem_seq
+            if len(sem_seqs.token)>0:
+                random_id_list = [k for k in range(len(sem_seqs.token))] 
+                random.shuffle(random_id_list)
+                for key,value in vars(sem_seqs).items():
+                    _, value = zip(*sorted(zip(random_id_list, value)))
+                    setattr(sem_seqs, key, list(value))
+                _, num_sem_list = zip(*sorted(zip(random_id_list, num_sem_list)))
+            # expression map 
+            for k in range(len(sem_seqs.token)):
+                num_map_list += num_sem_list[k]
+            num_map_dict = {key:value for key, value in zip(num_map_list, num_all_list)} 
+            for k in range(len(exp)):
+                if exp[k] in num_map_dict:
+                    exp[k] = num_map_dict[exp[k]]
+
+class StruSeq_RandomRotate(object):
+
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, text_seq, stru_seqs, sem_seqs, exp):
+        if random.random() < self.prob:
+            # shuffle stru_seq
+            if len(stru_seqs.token)>0:
+                random_id_list = [k for k in range(len(stru_seqs.token))] 
+                random.shuffle(random_id_list)
+                for key, value in vars(stru_seqs).items():
+                    _, value = zip(*sorted(zip(random_id_list, value)))
+                    setattr(stru_seqs, key, list(value))

datasets/utils.py

@@ -0,0 +1,266 @@
+punctuation_list = ['.', '?', ',']
+digit_list = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
+capital_letter_list = [chr(item) for item in range(65, 91)]
+low_letter_list = [chr(item) for item in range(97, 123)]
+begin_words = ["find", "what", "solve", "determine", "express", "how"]
+end_words = [".", ",", '?', "if", "so", "for which", "given", "with", "on",
+             "in", "must", 'for', 'that', 'formed']
+unit_list = ["mm^{2}", "cm^{2}", "in^{2}", "ft^{2}",
+             "yd^{2}", "km^{2}", "units^{2}", "mi^{2}", "m^{2}"]
+special_token_list = ['\\frac', '\\pi', '\\sqrt', "+", "-", "^"]
+
+
+def get_token(ss):
+    """
+        Tokenizer: divide the textual problem into words
+    """
+    raw_str_list = ss.strip().split(' ')
+    # Split punctuation
+    new_str1_list = []
+    for item in raw_str_list:
+        if item[-1] in punctuation_list:
+            new_str1_list.append(item[:-1])
+            new_str1_list.append(item[-1])
+        else:
+            new_str1_list.append(item)
+    # Split points (capital letters)
+    new_str2_list = []
+    for item in new_str1_list:
+        is_geo_rep = True
+        point_list = []
+        for k in item:
+            if (ord(k) >= 65 and ord(k) <= 90) or \
+                    ((k == '\'' or k in digit_list) and len(point_list) > 0):
+                if k == '\'' or k in digit_list:
+                    point_list[-1] += k
+                else:
+                    point_list.append(k)
+            else:
+                is_geo_rep = False
+                break
+        if is_geo_rep:
+            new_str2_list += point_list
+        else:
+            new_str2_list.append(item.lower())
+
+    return new_str2_list
+
+def split_text(text_data):
+    """
+        split textual problem into condition and problem(target)
+    """
+    if len(text_data.token) == 0:
+        return
+    begin_ind = 0
+    end_ind = len(text_data.token)
+    for id, token in enumerate(text_data.token):
+        if token in begin_words:
+            begin_ind = id
+            break
+    for id in range(begin_ind+2, len(text_data.token)):
+        if text_data.token[id] in end_words:
+            if text_data.token[id] in punctuation_list:
+                end_ind = id + 1
+            else:
+                end_ind = id
+            break
+    text_data.sect_tag = ['[COND]']*len(text_data.token[:begin_ind]) + \
+                            ['[PROB]']*len(text_data.token[begin_ind: end_ind]) + \
+                            ['[COND]']*len(text_data.token[end_ind:])
+
+def get_point_angleID_tag(text_data, stru_data, sem_data):
+    for id, item in enumerate(text_data.token):
+        if item[0] in capital_letter_list:
+           text_data.class_tag[id] = '[POINT]'
+        if item.isdigit() and id > 0 and text_data.token[id-1] == "\\angle":
+           text_data.class_tag[id] = '[ANGID]'
+
+    for k in range(len(stru_data.token)):
+        for id, item in enumerate(stru_data.token[k]):
+            if item[0] in capital_letter_list:
+                stru_data.class_tag[k][id] = '[POINT]'
+            if item.isdigit() and id > 0 and stru_data.token[k][id-1] == "\\angle":
+                stru_data.class_tag[k][id] = '[ANGID]'
+
+    for k in range(len(sem_data.token)):
+        for id, item in enumerate(sem_data.token[k]):
+            if item[0] in capital_letter_list:
+                sem_data.class_tag[k][id] = '[POINT]'
+            if item.isdigit() and id > 0 and sem_data.token[k][id-1] == "\\angle":
+                sem_data.class_tag[k][id] = '[ANGID]'
+
+def get_args(token):
+    letter_list = []
+    for special_token in special_token_list:
+        token = token.replace(special_token, "")
+    for letter in token:
+        if letter in low_letter_list and not letter in letter_list:
+            letter_list.append(letter)
+    return letter_list
+
+def get_num_arg_tag(text_data, sem_data):
+    """
+        Determine the variables/arguments in the text condition
+    """
+    arg_sem_flat = []
+    for k in range(len(sem_data.token)):
+        if len(sem_data.token[k]) >= 3 and sem_data.token[k][-3] == '=':
+            sem_data.class_tag[k][-2] = '[NUM]'
+            arg_sem_flat += get_args(sem_data.token[k][-2])
+
+    for id, token in enumerate(text_data.token):
+        if text_data.sect_tag[id] == '[COND]' and text_data.class_tag[id] == '[GEN]':
+            # unit symbol
+            if token in unit_list:
+                continue
+            # digit existing (rough judgment)
+            for word in digit_list:
+                if word in token:
+                    text_data.class_tag[id] = '[NUM]'
+                    break
+            # There are special characters, but not only special characters
+            for word in special_token_list:
+                if word in token and word != token:
+                    text_data.class_tag[id] = '[NUM]'
+                    break
+            # Single lowercase letter, but not special cases
+            if text_data.token[id] in low_letter_list:
+                if id < len(text_data.token)-1 and text_data.token[id+1] == '=':
+                    continue
+                if text_data.token[id] == 'm' and id < len(text_data.token)-1 and text_data.token[id+1] in ["\\angle", "\\widehat"]:
+                    continue
+                if text_data.token[id] == 'a' and (id == 0 or text_data.token[id-1] != '='):
+                    continue
+                if not text_data.token[id] in arg_sem_flat and \
+                    id > 0 and ('line' in text_data.token[id-1] or text_data.token[id-1] == 'and' or
+                                (text_data.token[id-1] == ',' and text_data.token[id+1] == ',')):
+                    continue
+                text_data.class_tag[id] = '[NUM]'
+
+    arg_text_flat = []
+    for id, token in enumerate(text_data.token):
+        if text_data.sect_tag[id] == '[COND]' and text_data.class_tag[id] == '[NUM]':
+            arg_text_flat += get_args(token)
+
+    # Determine arguments
+    arg_all_flat = arg_text_flat + arg_sem_flat
+    for id, token in enumerate(text_data.token):
+        if text_data.class_tag[id] == '[GEN]' \
+                and text_data.token[id] in arg_all_flat:
+            if id < len(text_data.token)-1 and text_data.token[id+1] == '=':
+                text_data.class_tag[id] = '[ARG]'
+                continue
+            if text_data.token[id] == 'm' and id < len(text_data.token)-1 and text_data.token[id+1] in ["\\angle", "\\widehat"]:
+                continue
+            if text_data.token[id] == 'a' and (id == 0 or text_data.token[id-1] != '=') and \
+                                text_data.sect_tag[id]=='[COND]':
+                continue
+            if id > 0 and ('line' in text_data.token[id-1] or text_data.token[id-1] == 'and' or
+                           (text_data.token[id-1] == ',' and text_data.token[id+1] == ',')):
+                continue
+            text_data.class_tag[id] = '[ARG]'
+
+def remove_sem_dup(text_data, sem_data, exp_token):
+    """
+        Remove the seq of sem_data if num is also in the text_data
+        and change the corresponding expression
+    """
+    text_num_list, id_all_list, id_map_list = [], [], []
+    token_, sect_tag_, class_tag_ = [], [], []
+
+    for k in range(len(text_data.token)):
+        if text_data.class_tag[k] == '[NUM]':
+            text_num_list.append(text_data.token[k])
+            var_name = 'N'+str(len(id_all_list))
+            id_all_list.append(var_name)
+            id_map_list.append(var_name)
+
+    for k in range(len(sem_data.token)):
+        if sem_data.class_tag[k][-2] == '[NUM]':
+            var_name = 'N'+str(len(id_all_list))
+            id_all_list.append(var_name)  
+            if not sem_data.token[k][-2] in text_num_list:
+                token_.append(sem_data.token[k])
+                sect_tag_.append(sem_data.sect_tag[k])
+                class_tag_.append(sem_data.class_tag[k])
+                id_map_list.append(var_name)
+        else:
+            token_.append(sem_data.token[k])
+            sect_tag_.append(sem_data.sect_tag[k])
+            class_tag_.append(sem_data.class_tag[k])
+
+    num_map_dict = {key:value for key, value in zip(id_map_list, id_all_list)} 
+    for k in range(len(exp_token)):
+        if exp_token[k] in num_map_dict:
+            exp_token[k] = num_map_dict[exp_token[k]]
+
+    sem_data.token = token_
+    sem_data.sect_tag = sect_tag_
+    sem_data.class_tag = class_tag_
+    
+def get_combined_text(text_seq, stru_seqs, sem_seqs, combine_text, args):
+    '''
+        combination style:  [stru_seqs, text_cond, sem_seqs, text_prob]
+    '''
+    # split cond and prob in text_seq
+    begin_ind = end_ind = None
+    for k in range(len(text_seq.sect_tag)):
+        if text_seq.sect_tag[k]=='[PROB]': 
+            begin_ind = k
+            break
+    for k in range(len(text_seq.sect_tag)-1,-1,-1):
+        if text_seq.sect_tag[k]=='[PROB]':
+            end_ind = k+1
+            break
+    # combine text_seq, stru_seqs and sem_seqs
+    for key in vars(combine_text):
+        # get text_cond and text_prob
+        text_all_value = getattr(text_seq, key)
+        text_cond_value = text_all_value[:begin_ind] + text_all_value[end_ind:]
+        text_prob_value = text_all_value[begin_ind:end_ind]
+        if args.without_stru:
+            value_all = text_cond_value + sum(getattr(sem_seqs, key), []) + text_prob_value
+        else:
+            value_all = sum(getattr(stru_seqs, key), []) + text_cond_value + \
+                                sum(getattr(sem_seqs, key), []) + text_prob_value
+                
+        setattr(combine_text, key, value_all)
+    
+def get_var_arg(combine_text, args):
+
+    var_values, arg_values = [], []
+    var_positions, arg_positions = [], []
+    class_tag = combine_text.class_tag
+    token = combine_text.token
+
+    for k in range(len(class_tag)):
+        if class_tag[k] == '[NUM]':
+            var_values.append(token[k])
+            var_positions.append(k)
+        if class_tag[k] == '[ARG]':
+            arg_values.append(token[k])
+            arg_positions.append(k)
+    # merge position of var and arg
+    return  var_positions+arg_positions, var_values, arg_values
+
+def get_text_index(combine_text, src_lang):
+    
+    text_sect_tag = src_lang.indexes_from_sentence(combine_text.sect_tag, id_type='sect_tag')
+    text_class_tag = src_lang.indexes_from_sentence(combine_text.class_tag, id_type='class_tag')
+    text_token = [combine_text.token[:], ['[PAD]']*len(combine_text.token)]
+    for k in range(len(combine_text.class_tag)):
+        if combine_text.class_tag[k] == '[NUM]':
+            letter_list = get_args(combine_text.token[k])
+            text_token[0][k] = text_token[1][k] = "[PAD]"
+            for j in range(len(letter_list)):
+                text_token[j][k] = letter_list[j]
+    text_token = [src_lang.indexes_from_sentence(item, id_type='text') for item in text_token]
+
+    return text_token, text_sect_tag, text_class_tag
+
+
+            
+
+
+
+    

loss/__init__.py

@@ -0,0 +1,10 @@
+from .loss import *
+from config import criterion_list
+
+
+def get_criterion(args):   
+    # create model
+    if args.criterion in criterion_list:
+        return eval(args.criterion)(args)
+    else:
+        raise NotImplementedError("Unsupported Loss Criterion : {}".format(args.criterion)) 

loss/loss.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from utils import *
+
+
+class CrossEntropy(nn.Module):
+    def __init__(self, cfg):
+        super(CrossEntropy, self).__init__()
+
+    def forward(self, output, target):
+        loss = F.cross_entropy(output, target)
+        return loss
+
+class FocalLoss(nn.Module):
+    def __init__(self, cfg=None):
+        super(FocalLoss, self).__init__()
+        # self.gamma = cfg.LOSS.FOCAL.GAMMA
+        if cfg is None:
+            self.gamma = 2.0
+        else:
+            self.gamma = cfg.focal_loss_gamma
+        assert self.gamma >= 0
+
+    def focal_loss(self, input_values):
+        """Computes the focal loss"""
+        p = torch.exp(-input_values)
+        loss = (1 - p) ** self.gamma * input_values
+        return loss.mean()
+
+    def forward(self, input, target):
+        return self.focal_loss(F.cross_entropy(input, target, reduction='none'))
+
+class MaskedCrossEntropy(nn.Module):
+
+    def __init__(self, cfg):
+        super(MaskedCrossEntropy, self).__init__()
+        self.cfg = cfg
+    
+    def forward(self, logits, target, length):
+        """
+        Args:
+            logits: A Variable containing a FloatTensor of size
+                (batch, max_len, num_classes) which contains the
+                unnormalized probability for each class.  B x S x (op_size+const_size+var_size)
+            target: A Variable containing a LongTensor of size
+                (batch, max_len) which contains the index of the true
+                class for each corresponding step. B x S
+        Returns:
+            loss: An average loss value masked by the length.
+        """
+        # logits_flat: (batch * max_len, num_classes)
+        logits_flat = logits.view(-1, logits.size(-1)) 
+        # log_probs_flat: (batch * max_len, num_classes)
+        log_probs_flat = F.log_softmax(logits_flat, dim=1)
+        # target_flat: (batch * max_len, 1)
+        target_flat = target.view(-1, 1)
+        # losses_flat: (batch * max_len, 1)
+        losses_flat = -torch.gather(log_probs_flat, dim=1, index=target_flat)
+        # losses: (batch, max_len)
+        losses = losses_flat.view(*target.size())
+        # mask: (batch, max_len)
+        mask = sequence_mask(length)
+        losses = losses * mask.float()
+        loss = losses.sum() / length.float().sum()
+        return loss

model/backbone/__init__.py

@@ -0,0 +1,16 @@
+from .resnet import *
+from .mobilenet_v2 import *
+from config import visual_backbone_list
+
+
+def get_visual_backbone(args):
+    if args.visual_backbone in visual_backbone_list:
+        model = eval(args.visual_backbone)()
+        if args.pretrain_vis_path !="":
+            model.load_model(pretrain=args.pretrain_vis_path)
+            args.logger.info("Visual backbone has been loaded...")
+        else:
+            args.logger.info("Visual backbone choose to train from scratch")
+        return model
+    else:
+        raise NotImplementedError("Unsupported Backbone: {}".format(args.visual_backbone))

model/backbone/mobilenet_v2.py

@@ -0,0 +1,122 @@
+import torch.nn as nn
+import math
+import torch
+import config as cfg
+
+
+def conv_bn(inp, oup, stride):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        nn.BatchNorm2d(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+            )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, width_mult=1.):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        interverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            [6, 160, 3, 2],
+            [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        # assert input_size % 32 == 0
+        input_channel = int(input_channel * width_mult)
+        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
+        self.features = [conv_bn(3, input_channel, 2)]
+        # building inverted residual blocks
+        for t, c, n, s in interverted_residual_setting:
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                if i == 0:
+                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
+                else:
+                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
+                input_channel = output_channel
+
+        # make it nn.Sequential
+        self.features = nn.Sequential(*self.features)
+
+        self._initialize_weights()
+
+    def forward(self, x):
+        x = self.features(x)
+        return x
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+                if m.bias is not None:
+                    m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                n = m.weight.size(1)
+                m.weight.data.normal_(0, 0.01)
+                m.bias.data.zero_()
+
+    def load_model(self):
+        model_dict = self.state_dict()
+        pretrained_dict = torch.load(cfg.pretrained_model_path)
+        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
+        model_dict.update(pretrained_dict)
+        self.load_state_dict(model_dict)
+
+
+def mobilenet_v2():
+
+    return MobileNetV2()

model/backbone/resnet.py

@@ -0,0 +1,159 @@
+import torch
+import torch.nn as nn
+import math
+
+def init_layer(L):
+    # Initialization using fan-in
+    if isinstance(L, nn.Conv2d):
+        n = L.kernel_size[0]*L.kernel_size[1]*L.out_channels
+        L.weight.data.normal_(0,math.sqrt(2.0/float(n)))
+    elif isinstance(L, nn.BatchNorm2d):
+        L.weight.data.fill_(1)
+        L.bias.data.fill_(0)
+
+class Flatten(nn.Module):
+    def __init__(self):
+        super(Flatten, self).__init__()
+        
+    def forward(self, x):        
+        return x.view(x.size(0), -1)
+
+# Simple ResNet Block
+class SimpleBlock(nn.Module):
+    maml = False #Default
+    def __init__(self, indim, outdim, half_res):
+        super(SimpleBlock, self).__init__()
+        self.indim = indim
+        self.outdim = outdim
+        self.C1 = nn.Conv2d(indim, outdim, kernel_size=3, stride=2 if half_res else 1, padding=1, bias=False)
+        self.BN1 = nn.BatchNorm2d(outdim)
+        self.C2 = nn.Conv2d(outdim, outdim,kernel_size=3, padding=1,bias=False)
+        self.BN2 = nn.BatchNorm2d(outdim)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.parametrized_layers = [self.C1, self.C2, self.BN1, self.BN2]
+
+        self.half_res = half_res
+
+        # if the input number of channels is not equal to the output, then need a 1x1 convolution
+        if indim!=outdim:
+            self.shortcut = nn.Conv2d(indim, outdim, 1, 2 if half_res else 1, bias=False)
+            self.BNshortcut = nn.BatchNorm2d(outdim)
+            self.parametrized_layers.append(self.shortcut)
+            self.parametrized_layers.append(self.BNshortcut)
+            self.shortcut_type = '1x1'
+        else:
+            self.shortcut_type = 'identity'
+
+        for layer in self.parametrized_layers:
+            init_layer(layer)
+
+    def forward(self, x):
+        out = self.C1(x)
+        out = self.BN1(out)
+        out = self.relu1(out)
+        out = self.C2(out)
+        out = self.BN2(out)
+        short_out = x if self.shortcut_type == 'identity' else self.BNshortcut(self.shortcut(x))
+        out = out + short_out
+        out = self.relu2(out)
+        return out
+
+# Bottleneck block
+class BottleneckBlock(nn.Module):
+    maml = False #Default
+    def __init__(self, indim, outdim, half_res):
+        super(BottleneckBlock, self).__init__()
+        bottleneckdim = int(outdim/4)
+        self.indim = indim
+        self.outdim = outdim
+
+        self.C1 = nn.Conv2d(indim, bottleneckdim, kernel_size=1,  bias=False)
+        self.BN1 = nn.BatchNorm2d(bottleneckdim)
+        self.C2 = nn.Conv2d(bottleneckdim, bottleneckdim, kernel_size=3, stride=2 if half_res else 1,padding=1)
+        self.BN2 = nn.BatchNorm2d(bottleneckdim)
+        self.C3 = nn.Conv2d(bottleneckdim, outdim, kernel_size=1, bias=False)
+        self.BN3 = nn.BatchNorm2d(outdim)
+
+        self.relu = nn.ReLU()
+        self.parametrized_layers = [self.C1, self.BN1, self.C2, self.BN2, self.C3, self.BN3]
+        self.half_res = half_res
+
+        # if the input number of channels is not equal to the output, then need a 1x1 convolution
+        if indim!=outdim:
+            self.shortcut = nn.Conv2d(indim, outdim, 1, stride=2 if half_res else 1, bias=False)
+            self.parametrized_layers.append(self.shortcut)
+            self.shortcut_type = '1x1'
+        else:
+            self.shortcut_type = 'identity'
+
+        for layer in self.parametrized_layers:
+            init_layer(layer)
+
+    def forward(self, x):
+
+        short_out = x if self.shortcut_type == 'identity' else self.shortcut(x)
+        out = self.C1(x)
+        out = self.BN1(out)
+        out = self.relu(out)
+        out = self.C2(out)
+        out = self.BN2(out)
+        out = self.relu(out)
+        out = self.C3(out)
+        out = self.BN3(out)
+        out = out + short_out
+
+        out = self.relu(out)
+        return out
+
+class ResNet(nn.Module):
+    maml = False #Default
+    def __init__(self,block,list_of_num_layers, list_of_out_dims, flatten = True):
+        # list_of_num_layers specifies number of layers in each stage
+        # list_of_out_dims specifies number of output channel for each stage
+        super(ResNet,self).__init__()
+        assert len(list_of_num_layers)==4, 'Can have only four stages'
+
+        conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                                            bias=False)
+        bn1 = nn.BatchNorm2d(64)
+        relu = nn.ReLU()
+        pool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        init_layer(conv1)
+        init_layer(bn1)
+        trunk = [conv1, bn1, relu, pool1]
+        indim = 64
+        for i in range(4):
+            for j in range(list_of_num_layers[i]):
+                half_res = (i>=1) and (j==0)
+                B = block(indim, list_of_out_dims[i], half_res)
+                trunk.append(B)
+                indim = list_of_out_dims[i]
+        if flatten:
+            avgpool = nn.AvgPool2d(4)
+            trunk.append(avgpool)
+            trunk.append(Flatten())
+            self.final_feat_dim = indim
+        else:
+            self.final_feat_dim = [indim, 4, 4]
+        self.trunk = nn.Sequential(*trunk)
+
+    def forward(self,x):
+        out = self.trunk(x)
+        return out
+
+def ResNet10(flatten = True):
+    return ResNet(SimpleBlock, [1,1,1,1],[64,128,256,512], flatten)
+
+def ResNet18(flatten = True):
+    return ResNet(SimpleBlock, [2,2,2,2],[64,128,256,512], flatten)
+
+def ResNet34(flatten = True):
+    return ResNet(SimpleBlock, [3,4,6,3],[64,128,256,512], flatten)
+
+def ResNet50(flatten = True):
+    return ResNet(BottleneckBlock, [3,4,6,3], [256,512,1024,2048], flatten)
+
+def ResNet101(flatten = True):
+    return ResNet(BottleneckBlock, [3,4,23,3],[256,512,1024,2048], flatten)

model/classifier/__init__.py

@@ -0,0 +1,23 @@
+from .classifier_ops import *
+from config import classifier_list
+
+       
+def get_classifier(args):
+
+    bias_flag = args.classifier_bias
+    num_features = args.num_features
+    num_classes = args.num_classes
+
+    if not args.classifier in classifier_list:
+        raise NotImplementedError("Unsupported Classifier: {}".format(args.classifier))
+
+    if args.classifier == "FCNorm":
+        classifier = FCNorm(num_features, num_classes)
+    elif args.classifier == "CosNorm":
+        classifier = CosNorm(num_features, num_classes)
+    elif args.classifier == "DotProduct":
+        classifier = DotProduct(num_classes, num_features, bias_flag)
+    elif args.classifier == "DistFC":
+        classifier = DistFC(num_features, num_classes)
+ 
+    return classifier

model/classifier/classifier_ops.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+class DotProduct(nn.Module):
+    def __init__(self, num_classes=1000, feat_dim=2048, bias=True):
+        super(DotProduct, self).__init__()
+        # print('<DotProductClassifier> contains bias: {}'.format(bias))
+        self.fc = nn.Linear(feat_dim, num_classes,bias)
+        
+    def forward(self, x, *args):
+        x = self.fc(x)
+        return x
+
+
+class CosNorm(nn.Module):
+    def __init__(self, in_dims, out_dims, scale=16, margin=0.5, init_std=0.001):
+        super(CosNorm, self).__init__()
+        self.in_dims = in_dims
+        self.out_dims = out_dims
+        self.scale = scale
+        self.margin = margin
+        self.weight = nn.Parameter(torch.Tensor(out_dims, in_dims).cuda())
+        self.reset_parameters() 
+
+    def reset_parameters(self):
+        stdv = 1. / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, *args):
+        norm_x = torch.norm(input.clone(), 2, 1, keepdim=True)
+        ex = (norm_x / (1 + norm_x)) * (input / norm_x)
+        ew = self.weight / torch.norm(self.weight, 2, 1, keepdim=True)
+        return torch.mm(self.scale * ex, ew.t())
+
+
+class FCNorm(nn.Module):
+    # for LDAM Loss
+    def __init__(self, num_features, num_classes, scale=20.0):
+        super(FCNorm, self).__init__()
+        self.weight = nn.Parameter(torch.FloatTensor(num_classes, num_features))
+        self.weight.data.uniform_(-1, 1).renorm_(2, 1, 1e-5).mul_(1e5)
+        self.scale = scale
+
+    def forward(self, x):
+        out = self.scale * F.linear(F.normalize(x), F.normalize(self.weight))
+        return out
+
+
+class DistFC(nn.Module):
+
+    def __init__(self, num_features, num_classes,init_weight=True):
+        super(DistFC, self).__init__()
+        self.centers=nn.Parameter(torch.randn(num_features,num_classes).cuda(),requires_grad=True)
+        if init_weight:
+            self.__init_weight()
+
+    def __init_weight(self):
+        nn.init.kaiming_normal_(self.centers)
+
+    def forward(self, x):
+        features_square=torch.sum(torch.pow(x,2),1, keepdim=True)
+        centers_square=torch.sum(torch.pow(self.centers,2),0, keepdim=True)
+        features_into_centers=2.0*torch.matmul(x, (self.centers))
+        dist=features_square+centers_square-features_into_centers   
+        return self.centers, dist
+
+

model/decoder/__init__.py

@@ -0,0 +1,24 @@
+# from .transformer import TransformerModel
+from config import decoder_list
+from .rnn_decoder import DecoderRNN
+from .tree_decoder import TreeDecoder
+from .transformer import TransformerDecoder
+
+def get_decoder(params, *args):
+         
+    if not params.decoder_type in decoder_list:
+        raise NotImplementedError(
+            "Unsupported Classifier: {}".format(params.decoder_type))
+
+    if params.decoder_type == "transformer":
+        decoder = TransformerDecoder(params, *args)
+    elif params.decoder_type == "rnn_decoder":
+        decoder = DecoderRNN(params, *args)
+    elif params.decoder_type == "tree_decoder":
+        decoder = TreeDecoder(params, *args)
+    else:
+        raise NotImplementedError("Unsupported Decoder: {}".format(params.decoder_type))
+             
+    return decoder
+
+

model/decoder/rnn_decoder.py

@@ -0,0 +1,201 @@
+import torch
+import torch.nn as nn
+from model.module import *
+from utils import *
+from torch.nn import functional as F
+
+class DecoderRNN(nn.Module):
+    def __init__(self, cfg, tgt_lang):
+        super(DecoderRNN, self).__init__()
+        # token location
+        self.var_start = tgt_lang.var_start # spe_num + midvar_num + const_num + op_num
+        self.sos_id = tgt_lang.word2index["[SOS]"]
+        self.eos_id = tgt_lang.word2index["[EOS]"]
+        # Define layers
+        self.em_dropout = nn.Dropout(cfg.dropout_rate)
+        self.embedding_tgt = nn.Embedding(self.var_start, cfg.decoder_embedding_size, padding_idx=0)
+        self.gru = nn.GRU(input_size=cfg.decoder_hidden_size+cfg.decoder_embedding_size, \
+                            hidden_size=cfg.decoder_hidden_size, \
+                            num_layers=cfg.decoder_layers, \
+                            dropout = cfg.dropout_rate, \
+                            batch_first = True)
+        # Choose attention model
+        self.attn = Attn(cfg.encoder_hidden_size, cfg.decoder_hidden_size)
+        self.score = Score(cfg.encoder_hidden_size+cfg.decoder_hidden_size, cfg.decoder_embedding_size)
+        # predefined constant
+        self.no_var_id = torch.arange(self.var_start).unsqueeze(0).cuda()
+        self.cfg = cfg
+    
+    def get_var_encoder_outputs(self, encoder_outputs, var_pos):
+        """
+        Arguments:
+            encoder_outputs:  B x S1 x H
+            var_pos: B x S3
+        Returns:
+            var_embeddings: B x S3 x H
+        """
+        hidden_size = encoder_outputs.size(-1)
+        expand_var_pos = var_pos.unsqueeze(-1).repeat(1, 1, hidden_size)
+        var_embeddings = encoder_outputs.gather(dim=1, index = expand_var_pos)
+        return var_embeddings
+    
+    def forward(self, encoder_outputs, problem_output, len_src, var_pos, len_var, \
+                            text_tgt=None, is_train=False):
+        """
+        Arguments:
+            encoder_outputs: B x S1 x H
+            problem_output: layer_num x B x H
+            len_src: B
+            text_tgt: B x S2
+            var_pos: B x S3
+            len_var: B
+        Return:
+            training: logits, B x S x (no_var_size+var_size)
+            testing: exp_id, B x candi_size(beam_size) x exp_len
+        """
+        self.embedding_var = self.get_var_encoder_outputs(encoder_outputs, var_pos) # B x S3 x H
+        self.src_mask = sequence_mask(len_src)  # B x S1
+        self.candi_mask = sequence_mask(self.var_start + len_var) # B x (no_var_size + var_size)
+        if is_train:
+            return self._forward_train(encoder_outputs, problem_output, text_tgt)  
+        else:
+            return self._forward_test(encoder_outputs, problem_output)
+            
+    def _forward_train(self, encoder_outputs, problem_output, text_tgt):
+
+        all_seq_outputs = []
+        batch_size = encoder_outputs.size(0)
+        # initial hidden input of RNN
+        rnn_hidden = problem_output
+        # input embedding
+        tgt_novar_id = torch.clamp(text_tgt, max=self.var_start-1) # B x S2
+        novar_embedding = self.embedding_tgt(tgt_novar_id) # B x S2 x H
+        tgt_var_id = torch.clamp(text_tgt-self.var_start, min=0) # B x S2
+        var_embeddings = self.embedding_var.gather(dim=1, index = \
+                            tgt_var_id.unsqueeze(2).repeat(1, 1, self.cfg.decoder_embedding_size)) # B x S2 x H
+
+        choose_mask = (text_tgt<self.var_start).unsqueeze(2). \
+                                repeat(1, 1, self.cfg.decoder_embedding_size)
+        embedding_all = torch.where(choose_mask, novar_embedding, var_embeddings) # B x S2 x H
+        embedding_all_ = self.em_dropout(embedding_all)
+        # candi weight embedding
+        embedding_weight_no_var = self.embedding_tgt(self.no_var_id. \
+                                    repeat(batch_size, 1)) # B x no_var_size x H
+        embedding_weight_all = torch.cat((embedding_weight_no_var, self.embedding_var), dim=1) # B x (no_var_size + var_size) x H
+        embedding_weight_all_ = self.em_dropout(embedding_weight_all) 
+
+        for t in range(text_tgt.size(1)-1):
+            # Calculate attention from current RNN state and all encoder outputs;
+            # apply to encoder outputs to get weighted average
+            current_hiddens = self.em_dropout(rnn_hidden[-1].unsqueeze(1)) # B x 1 x H
+            attn_weights = self.attn(current_hiddens, encoder_outputs, self.src_mask)
+            context = attn_weights.unsqueeze(1).bmm(encoder_outputs)  # B x 1 x H 
+            # Get current hidden state from input word and last hidden state
+            rnn_output, rnn_hidden = self.gru(torch.cat((embedding_all_[:, t:t+1, :], context), 2), rnn_hidden) 
+            # rnn_output: B x 1 x H 
+            # rnn_hidden: num_layers x B x H 
+            current_fusion_emb = torch.cat((rnn_output, context), 2)
+            current_fusion_emb_ = self.em_dropout(current_fusion_emb)
+            candi_score = self.score(current_fusion_emb_, embedding_weight_all_, \
+                                            self.candi_mask) #  B x (no_var_size + var_size)
+            all_seq_outputs.append(candi_score)
+        
+        all_seq_outputs = torch.stack(all_seq_outputs, dim=1)  
+
+        return all_seq_outputs
+    
+    def _forward_test(self, encoder_outputs, problem_output):
+        """
+            Decode with beam search algorithm
+        """
+        exp_outputs = []
+        batch_size = encoder_outputs.size(0)
+
+        for sample_id in range(batch_size):
+            # predefine 
+            rem_size = self.cfg.beam_size
+            encoder_output = encoder_outputs[sample_id:sample_id+1].repeat(rem_size, 1, 1) # beam_size x S1 x H
+            src_mask = self.src_mask[sample_id:sample_id+1].repeat(rem_size, 1) # beam_size x S1
+            embedding_var = self.embedding_var[sample_id:sample_id+1].repeat(rem_size, 1, 1) # beam_size x S3 x H
+            embedding_weight_no_var = self.embedding_tgt(self.no_var_id.repeat(rem_size, 1)) # beam_size x no_var_size x H
+            embedding_weight_all = torch.cat((embedding_weight_no_var, embedding_var), dim=1)  # beam_size x (no_var_size + var_size) x H
+            embedding_weight_all_ = self.em_dropout(embedding_weight_all) 
+            candi_mask = self.candi_mask[sample_id:sample_id+1].repeat(rem_size, 1) # beam_size x S1
+            candi_exp_output = []
+            candi_score_output = []
+            
+            for i in range(self.cfg.max_output_len):
+                # initial varible
+                if i==0:
+                    input_token = torch.LongTensor([[self.sos_id]]*rem_size).cuda() # rem_size x 1
+                    rnn_hidden = problem_output[:, sample_id:sample_id+1].repeat(1, rem_size, 1) # layer_num x rem_size x H
+                    current_score = torch.FloatTensor([[0.0]]*rem_size).cuda() # rem_size x 1
+                    current_exp_list = [[]]*rem_size
+                else:
+                    input_token = torch.LongTensor(token_list).unsqueeze(1).cuda() 
+                    rnn_hidden = rnn_hidden[:, cand_list]
+                    rem_size = len(exp_list)
+                    current_score = torch.FloatTensor(score_list[:rem_size]).unsqueeze(1).cuda()
+                    current_exp_list = exp_list
+                    
+                # input embedding
+                tgt_novar_id = torch.clamp(input_token, max=self.var_start-1) # rem_size x 1
+                novar_embedding = self.embedding_tgt(tgt_novar_id) # rem_size x 1 x H
+                tgt_var_id = torch.clamp(input_token-self.var_start, min=0) # rem_size x 1
+                var_embeddings = embedding_var[:rem_size].gather(dim=1, index=tgt_var_id.unsqueeze(2). \
+                                            repeat(1, 1, self.cfg.decoder_embedding_size)) # rem_size x 1 x H
+                choose_mask = (input_token<self.var_start).unsqueeze(2). \
+                                        repeat(1, 1, self.cfg.decoder_embedding_size) # rem_size x 1 x H
+                embedding_all = torch.where(choose_mask, novar_embedding, var_embeddings) # rem_size x 1 x H
+                embedding_all_ = self.em_dropout(embedding_all)
+                # attention 
+                current_hiddens = self.em_dropout(rnn_hidden[-1].unsqueeze(1))  # rem_size x 1 x H
+                attn_weights = self.attn(current_hiddens, encoder_output[:rem_size], src_mask[:rem_size]) # rem_size x S1 
+                context = attn_weights.unsqueeze(1).bmm(encoder_output[:rem_size])  # rem_size x 1 x H 
+                # Get current hidden state from input word and last hidden state
+                rnn_output, rnn_hidden = self.gru(torch.cat((embedding_all_, context), 2), rnn_hidden)
+                # rnn_output: rem_size x 1 x H 
+                # rnn_hidden: num_layers x rem_size x H 
+                current_fusion_emb = torch.cat((rnn_output, context), 2)
+                current_fusion_emb_ = self.em_dropout(current_fusion_emb)
+                candi_score = self.score(current_fusion_emb_, embedding_weight_all_[:rem_size], \
+                                                candi_mask[:rem_size]) #  rem_size x (no_var_size + var_size)
+                
+                if i==0:
+                    new_score = F.log_softmax(candi_score, dim=1)[:1]
+                else:
+                    new_score = F.log_softmax(candi_score, dim=1) + current_score
+
+                cand_tup_list = [(score, id) for id, score in enumerate(new_score.view(-1).tolist())]
+                cand_tup_list += [(score, -1) for score in candi_score_output]
+                cand_tup_list.sort(key=lambda x:x[0], reverse=True)
+
+                token_list = []
+                cand_list = []
+                exp_list = []
+                score_list = []
+
+                for tv, ti in cand_tup_list[:self.cfg.beam_size]:
+                    if ti!=-1:
+                      idex = ti
+                      x = idex // candi_score.size(-1) 
+                      y = idex % candi_score.size(-1)
+                      if y!=self.eos_id:
+                          token_list.append(y)
+                          cand_list.append(x)
+                          exp_list.append(current_exp_list[x]+[y])
+                          score_list.append(tv)
+                      else:
+                          candi_exp_output.append(current_exp_list[x])
+                          candi_score_output.append(float(tv))
+
+                if len(token_list)==0:
+                    break
+
+            if len(candi_exp_output)>0:
+                _, candi_exp_output = zip(*sorted(zip(candi_score_output, candi_exp_output), reverse=True))
+                exp_outputs.append(list(candi_exp_output[:self.cfg.beam_size]))
+            else:
+                exp_outputs.append([])
+
+        return exp_outputs

model/decoder/transformer.py

@@ -0,0 +1,217 @@
+import torch
+import torch.nn as nn
+from utils.utils import sequence_mask
+from model.module import *
+from torch.nn import functional as F
+import math
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, max_len=5000, dropout_rate=0.2):
+        super(PositionalEncoding, self).__init__()
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        """
+            x: [B, max_len, d_model]
+            pe: [1, max_len, d_model]
+        """
+        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
+        return self.dropout(x)
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, cfg, tgt_lang, \
+                d_model=256, nhead=8, num_decoder_layers=4, dim_feedforward=1024, dropout=0.2):
+        super(TransformerDecoder, self).__init__()
+
+        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
+        decoder_norm = nn.LayerNorm(d_model)
+        self.decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
+        self.position_dec = PositionalEncoding(d_model=d_model)
+
+        self.score = Score_Multi(cfg.decoder_hidden_size, cfg.decoder_embedding_size)
+        self.var_start = tgt_lang.var_start
+        self.embedding_tgt = nn.Embedding(self.var_start, cfg.decoder_embedding_size, padding_idx=0)
+        self.no_var_id = torch.arange(self.var_start).unsqueeze(0).cuda()
+
+        self._reset_parameters()
+        self.d_model = d_model
+        self.nhead = nhead
+        self.cfg = cfg
+        self.sos_id = tgt_lang.word2index["[SOS]"]
+        self.eos_id = tgt_lang.word2index["[EOS]"]
+        
+    def _reset_parameters(self):
+        """
+            Initiate parameters in the transformer model.
+        """
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def get_square_subsequent_mask(self, sz):
+        """
+            Generate a square mask for the sequence. The masked positions are filled with True.
+            Unmasked positions are filled with False.
+        """
+        mask = (torch.triu(torch.ones(sz, sz)) == 0).transpose(0, 1)
+        return mask.cuda()
+    
+    def get_var_encoder_outputs(self, encoder_outputs, var_pos):
+        """
+        Arguments:
+            encoder_outputs:  B x S1 x H
+            var_pos: B x S3
+        Returns:
+            var_embeddings: B x S3 x H
+        """
+        hidden_size = encoder_outputs.size(-1)
+        expand_var_pos = var_pos.unsqueeze(-1).repeat(1, 1, hidden_size)
+        var_embeddings = encoder_outputs.gather(dim=1, index = expand_var_pos)
+        return var_embeddings
+    
+    def forward(self, memory, len_src, tgt, len_tgt, var_pos, len_var, is_train=False):
+        '''
+            memory: B x S1 x H
+            len_src: B
+            tgt: B x S2
+            len_tgt: B
+            var_pos: B x S3(var_size)
+            len_var: B
+        '''
+        self.embedding_var = self.get_var_encoder_outputs(memory, var_pos) # B x S3 x H
+        self.candi_mask = sequence_mask(self.var_start + len_var) # B x (no_var_size + var_size)
+        self.memory_key_padding_mask = ~sequence_mask(len_src) # B x S1
+        if is_train:
+            return self._forward_train(memory, tgt, len_tgt)  
+        else:
+            return self._forward_test(memory)
+
+
+    def _forward_train(self, memory, tgt, len_tgt):
+        # mask
+        tgt_mask = self.get_square_subsequent_mask(tgt.size(-1))
+        tgt_key_padding_mask = ~sequence_mask(len_tgt)
+        # emb_tgt
+        tgt_novar_id = torch.clamp(tgt, max=self.var_start-1) # B x S2
+        novar_embedding = self.embedding_tgt(tgt_novar_id) # B x S2 x H
+        tgt_var_id = torch.clamp(tgt-self.var_start, min=0) # B x S2
+        var_embeddings = self.embedding_var.gather(dim=1, index = \
+                            tgt_var_id.unsqueeze(2).repeat(1, 1, self.cfg.decoder_embedding_size)) # B x S2 x H
+        choose_mask = (tgt<self.var_start).unsqueeze(2). \
+                                repeat(1, 1, self.cfg.decoder_embedding_size)
+        emb_tgt = torch.where(choose_mask, novar_embedding, var_embeddings) # B x S2 x H
+        # position decoding
+        emb_tgt = self.position_dec(emb_tgt)
+        output = self.decoder( # B x S2 x H
+            emb_tgt.permute(1,0,2), 
+            memory.permute(1,0,2), 
+            tgt_mask=tgt_mask, 
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=self.memory_key_padding_mask,
+        ).permute(1,0,2)
+        # candi weight embedding
+        embedding_weight_no_var = self.embedding_tgt(self.no_var_id.repeat(len(len_tgt), 1)) # B x no_var_size x H
+        embedding_weight_all = torch.cat((embedding_weight_no_var, self.embedding_var), dim=1)  # B x (no_var_size+var_size) x H
+        candi_score = self.score( #  B x S2 x (no_var_size + var_size)
+            output, 
+            embedding_weight_all, \
+            self.candi_mask
+            ) 
+
+        return candi_score[:,:-1,:].clone()
+
+    def _forward_test(self, memory):
+        
+        exp_outputs = []
+
+        for sample_id in range(memory.size(0)):
+            # predefine 
+            rem_size = self.cfg.beam_size
+            memory_item = memory[sample_id:sample_id+1].repeat(rem_size, 1, 1) # beam_size x S1 x H
+            memory_key_padding_mask = self.memory_key_padding_mask[sample_id:sample_id+1].repeat(rem_size, 1) # beam_size x S1
+            embedding_var = self.embedding_var[sample_id:sample_id+1].repeat(rem_size, 1, 1) # beam_size x S3 x H
+            embedding_weight_no_var = self.embedding_tgt(self.no_var_id.repeat(rem_size, 1)) # beam_size x no_var_size x H
+            embedding_weight_all = torch.cat((embedding_weight_no_var, embedding_var), dim=1)  # beam_size x (no_var_size + var_size) x H
+            candi_mask = self.candi_mask[sample_id:sample_id+1].repeat(rem_size, 1) # beam_size x S1
+
+            candi_exp_output = []
+            candi_score_output = []
+            
+            tgt = torch.LongTensor([[self.sos_id]]*rem_size).cuda() # rem_size x 1
+            len_tgt = torch.LongTensor([1]*rem_size).cuda() # rem_size
+            current_score = torch.FloatTensor([[0.0]]*rem_size).cuda() # rem_size x 1
+            current_exp_list = [[self.sos_id]]*rem_size
+
+            for i in range(self.cfg.max_output_len):
+                # mask
+                tgt_mask = self.get_square_subsequent_mask(tgt.size(-1))
+                tgt_key_padding_mask = ~sequence_mask(len_tgt)
+                # input embedding
+                tgt_novar_id = torch.clamp(tgt, max=self.var_start-1) # rem_size x S
+                novar_embedding = self.embedding_tgt(tgt_novar_id) # rem_size x S x H
+                tgt_var_id = torch.clamp(tgt-self.var_start, min=0) # rem_size x S
+                var_embeddings = embedding_var[:rem_size].gather(dim=1, index=tgt_var_id.unsqueeze(2). \
+                                                    repeat(1, 1, self.cfg.decoder_embedding_size)) # rem_size x S x H
+                choose_mask = (tgt<self.var_start).unsqueeze(2).repeat(1, 1, self.cfg.decoder_embedding_size) # rem_size x S x H
+                emb_tgt = torch.where(choose_mask, novar_embedding, var_embeddings) # rem_size x S x H
+                # position decoding
+                emb_tgt = self.position_dec(emb_tgt)
+                output = self.decoder( # rem_size x S x H
+                    emb_tgt.permute(1,0,2), 
+                    memory_item[:rem_size].permute(1,0,2), 
+                    tgt_mask=tgt_mask, 
+                    tgt_key_padding_mask=tgt_key_padding_mask,
+                    memory_key_padding_mask=memory_key_padding_mask[:rem_size],
+                ).permute(1,0,2)
+                candi_score = self.score( # rem_size x S x (no_var_size + var_size)
+                    output, 
+                    embedding_weight_all[:rem_size], \
+                    candi_mask[:rem_size]
+                    ) 
+
+                if i==0:
+                    new_score = F.log_softmax(candi_score[:, -1, :], dim=1)[:1]
+                else:
+                    new_score = F.log_softmax(candi_score[:, -1, :], dim=1) + current_score # rem_size x (no_var_size + var_size)
+                
+                topv, topi = new_score.view(-1).topk(rem_size)
+                exp_list = []
+                score_list = topv.tolist()
+
+                for tv, ti in zip(topv, topi):
+                    idex = ti.item()
+                    x = idex // candi_score.size(-1) 
+                    y = idex % candi_score.size(-1)
+                    if y!=self.eos_id:
+                        exp_list.append(current_exp_list[x]+[y])
+                    else:
+                        candi_exp_output.append(current_exp_list[x][1:])
+                        candi_score_output.append(float(tv))
+
+                if len(exp_list)==0:
+                    break
+
+                tgt = torch.LongTensor(exp_list).cuda() # rem_size x S
+                len_tgt = torch.LongTensor([len(item) for item in exp_list]).cuda() # rem_size
+                current_exp_list = exp_list
+                rem_size = len(exp_list)
+                current_score = torch.FloatTensor(score_list[:rem_size]).unsqueeze(1).cuda() # rem_size x 1
+  
+            if len(candi_exp_output)>0:
+                _, candi_exp_output = zip(*sorted(zip(candi_score_output, candi_exp_output), reverse=True))
+                exp_outputs.append(list(candi_exp_output))
+            else:
+                exp_outputs.append([])
+
+        return exp_outputs

model/decoder/tree_decoder.py

@@ -0,0 +1,374 @@
+import torch
+import torch.nn as nn
+from utils import *
+from model.module import *
+from torch.nn import functional as F
+import copy
+
+class TreeNode:  # the class save the tree node
+    def __init__(self, embedding, left_flag=False):
+        self.embedding = embedding
+        self.left_flag = left_flag
+
+class TreeEmbedding:  # the class save the tree
+    def __init__(self, embedding, terminal=False):
+        self.embedding = embedding
+        self.terminal = terminal
+
+class TreeBeam:  # the class save the beam node
+    def __init__(self, score, node_stacks, embeddings_stacks, left_child_trees, out):
+        self.score = score
+        self.embeddings_stacks = embeddings_stacks
+        self.node_stacks = node_stacks
+        self.left_child_trees = left_child_trees
+        self.out = out
+
+class Prediction(nn.Module):
+    # a seq2tree decoder with Problem aware dynamic encoding
+    def __init__(self, cfg, op_const_size):
+        super(Prediction, self).__init__()
+        # Define layers
+        self.em_dropout = nn.Dropout(cfg.dropout_rate)
+        # for Computational symbols and Generated numbers
+        self.concat_l = nn.Linear(cfg.decoder_hidden_size, cfg.decoder_hidden_size)
+        self.concat_r = nn.Linear(cfg.decoder_hidden_size * 2, cfg.decoder_hidden_size)
+        self.concat_lg = nn.Linear(cfg.decoder_hidden_size, cfg.decoder_hidden_size)
+        self.concat_rg = nn.Linear(cfg.decoder_hidden_size * 2, cfg.decoder_hidden_size)
+        # attention module
+        self.attn = Attn(cfg.encoder_hidden_size, cfg.decoder_hidden_size)
+        self.score = Score(cfg.encoder_hidden_size+cfg.decoder_hidden_size, cfg.decoder_embedding_size)
+        # predefined constant
+        self.op_const_id = torch.arange(op_const_size).unsqueeze(0).cuda()
+        self.padding_hidden = torch.zeros(1, cfg.decoder_hidden_size).cuda()
+
+    def forward(self, node_stacks, left_child_trees, encoder_outputs, var_pades, source_mask, candi_mask, embedding_op_const):
+        '''
+        Augments:
+            node_stacks: [[TreeNode(_)]]*B, store the variable h
+            left_child_trees: [t]*B, store the representation of left tree
+            encoder_outputs: [B, S1, H]
+            var_pades: [B, S2, H], all_vars_encoder_outputs 
+            padding_hidden: [1, H]
+            source_mask: [B, S1], mask for source seq 
+            candi_mask: [B, op_size+const_size+var_size], mask for target seq
+        Returns:
+            num_score: [B x (op_size+const_size+var_size)]
+            current_embeddings: q [B x 1 x H], the target vector of the current node 
+            current_context: c [B x 1 x H], the context vector of the current node, is calculated using the target vector and encoder_outputs
+            current_all_embeddings: [B x (op_size+const_size+var_size) x H] e (M_op, M_con, h_loc^p) 
+        '''
+        current_embeddings = []
+
+        for node_list in node_stacks:
+            if len(node_list) == 0:
+                current_embeddings.append(self.padding_hidden)
+            else:
+                current_node = node_list[-1]
+                current_embeddings.append(current_node.embedding)
+
+        current_node_temp = [] # B x (1 x H)
+
+        for l, c in zip(left_child_trees, current_embeddings):
+            if l is None:
+                cd = self.em_dropout(c)
+                g = torch.tanh(self.concat_l(cd))
+                t = torch.sigmoid(self.concat_lg(cd))
+                current_node_temp.append(g*t)
+            else:
+                ld = self.em_dropout(l)
+                cd = self.em_dropout(c)
+                g = torch.tanh(self.concat_r(torch.cat((ld, cd), 1)))
+                t = torch.sigmoid(self.concat_rg(torch.cat((ld, cd), 1)))
+                current_node_temp.append(g*t) 
+
+        current_node = torch.stack(current_node_temp, dim=0) # B x 1 x H (q)
+        current_embeddings = self.em_dropout(current_node) 
+        current_attn = self.attn(current_embeddings, encoder_outputs, source_mask) # B x S
+        current_context = current_attn.unsqueeze(1).bmm(encoder_outputs)  # B x 1 x H (c)
+        leaf_input = torch.cat((current_node, current_context), 2) # B x 1 x 2H
+        
+        embedding_weight_op_const = embedding_op_const(self.op_const_id.repeat(var_pades.size(0), 1)) # B x var_size x H
+        embedding_weight_all = torch.cat((embedding_weight_op_const, var_pades), dim=1)  # B x (op_size+const_size+var_size) x H
+
+        leaf_input = self.em_dropout(leaf_input)
+        embedding_weight_all_ = self.em_dropout(embedding_weight_all)   
+        num_score = self.score(leaf_input, embedding_weight_all_, candi_mask) # B x (op_size+const_size+var_size)
+
+        return num_score, current_node, current_context, embedding_weight_all
+
+class GenerateNode(nn.Module):
+    def __init__(self, cfg, op_size):
+        super(GenerateNode, self).__init__()
+
+        self.embedding_size = cfg.decoder_embedding_size
+        self.hidden_size = cfg.decoder_hidden_size
+        self.op_size = op_size
+
+        self.em_dropout = nn.Dropout(cfg.dropout_rate)
+        self.generate_l = nn.Linear(self.hidden_size * 2 + self.embedding_size, self.hidden_size)
+        self.generate_r = nn.Linear(self.hidden_size * 2 + self.embedding_size, self.hidden_size)
+        self.generate_lg = nn.Linear(self.hidden_size * 2 + self.embedding_size, self.hidden_size)
+        self.generate_rg = nn.Linear(self.hidden_size * 2 + self.embedding_size, self.hidden_size)
+
+    def forward(self, current_embedding, node_label, current_context, embedding_op_const):
+        """
+            Generate the hidden node hl and hr of tree, according to the front part of eq(10)(11)
+        Arguments:
+            current_embedding: [B x 1 x H (q)], the target vector of the current node
+            node_label: [B (id)]
+            current_context: [B x 1 x H (c)], context vector of current node
+            embedding_op_const: Embedding of op_const
+        Returns:
+            left_child: [B x H (h)]
+            right_child: [B x H (h)]
+            token_embedding: [B x H (e(y|P) of op)]
+        """
+        node_label_op = torch.clamp(node_label, max=self.op_size-1)
+        current_embedding_ = self.em_dropout(current_embedding.squeeze(1))
+        current_context_ = self.em_dropout(current_context.squeeze(1))
+        token_embedding = embedding_op_const(node_label_op)
+        token_embedding_ = self.em_dropout(token_embedding)
+
+        l_child = torch.tanh(self.generate_l(torch.cat((current_embedding_, current_context_, token_embedding_), 1)))
+        l_child_g = torch.sigmoid(self.generate_lg(torch.cat((current_embedding_, current_context_, token_embedding_), 1)))
+        r_child = torch.tanh(self.generate_r(torch.cat((current_embedding_, current_context_, token_embedding_), 1)))
+        r_child_g = torch.sigmoid(self.generate_rg(torch.cat((current_embedding_, current_context_, token_embedding_), 1)))
+        l_child = l_child * l_child_g
+        r_child = r_child * r_child_g
+
+        return l_child, r_child, token_embedding
+
+class Merge(nn.Module):
+    """
+        Get subtree embedding via Recursive Neural Network
+    """
+    def __init__(self, cfg):
+        super(Merge, self).__init__()
+
+        self.embedding_size = cfg.decoder_embedding_size
+        self.hidden_size = cfg.decoder_hidden_size
+        self.em_dropout = nn.Dropout(cfg.dropout_rate)
+        self.merge = nn.Linear(self.hidden_size * 2 + self.embedding_size, self.hidden_size)
+        self.merge_g = nn.Linear(self.hidden_size * 2 + self.embedding_size, self.hidden_size)
+
+    def forward(self, node_embedding, sub_tree_1, sub_tree_2):
+        '''
+        Arguments:
+            node_embedding: 1 x H
+            sub_tree_1: 1 x H 
+            sub_tree_2: 1 x H
+        Return:
+            sub_tree: 1 x H
+        '''
+        sub_tree_1 = self.em_dropout(sub_tree_1)
+        sub_tree_2 = self.em_dropout(sub_tree_2)
+        node_embedding = self.em_dropout(node_embedding)
+
+        sub_tree = torch.tanh(self.merge(torch.cat((node_embedding, sub_tree_1, sub_tree_2), 1)))
+        sub_tree_g = torch.sigmoid(self.merge_g(torch.cat((node_embedding, sub_tree_1, sub_tree_2), 1)))
+        sub_tree = sub_tree * sub_tree_g
+        
+        return sub_tree
+
+class TreeDecoder(nn.Module):
+    def __init__(self, cfg, tgt_lang):
+        super(TreeDecoder, self).__init__()
+        # embedding for op, const, num
+        self.var_start = tgt_lang.var_start
+        self.op_num = tgt_lang.op_num
+        self.const_num = tgt_lang.const_num
+
+        self.embedding_op_const = nn.Embedding(self.op_num+self.const_num, cfg.decoder_embedding_size)
+        self.embedding_var = None # obtain from encoder
+        self.cfg = cfg
+        # modules of TreeDecoder 
+        self.predict = Prediction(cfg, self.op_num+self.const_num)
+        self.generate = GenerateNode(cfg, self.op_num)
+        self.merge = Merge(cfg)
+    
+    def get_var_encoder_outputs(self, encoder_outputs, var_positions):
+        """
+        Arguments:
+            encoder_outputs:  B x S1 x H
+            var_positions: B x S2
+        Returns:
+            var_embeddings: B x S2 x H
+        """
+        hidden_size = encoder_outputs.size(-1)
+        expand_var_positions = var_positions.unsqueeze(-1).repeat(1, 1, hidden_size)
+        var_embeddings = encoder_outputs.gather(dim=1, index = expand_var_positions)
+        return var_embeddings
+
+    def forward(self, encoder_outputs, problem_output, len_source, var_positions, len_var, \
+                            is_train=False, text_target=None, len_target=None):
+        """
+        Arguments:
+            encoder_outputs: B x S1 x H
+            problem_output: B x H
+            len_source: B
+            text_target: B x S2
+            len_target: B
+            var_positions: B x S3
+            len_var: B
+        Return:
+            training: output B x S x (op_size+const_size+var_size), logits of one batch
+            testing: [expr] x B
+        """
+        self.embedding_var = self.get_var_encoder_outputs(encoder_outputs, var_positions) # B x S2 x H
+        self.source_mask = sequence_mask(len_source)
+        self.candi_mask = sequence_mask(len_var+self.var_start)
+        if is_train:
+            return self._forward_train(encoder_outputs, problem_output, text_target)  
+        else:
+            return self._forward_test(encoder_outputs, problem_output)
+            
+    def _forward_train(self, encoder_outputs, problem_output, text_target):
+        """
+        Arguments:
+            embeddings_stacks: [[TreeEmbedding(t, terminal)]]*B, a stack of subtrees t in the first order traversal
+            left_child_trees: [t]*B, the representation of left tree of current node
+            node_stacks: [[TreeNode(h, left_flag)]]*B, a stack of hidden state h in the first order traversal
+        Returns:
+            all_node_outputs: B x S x (op_size+const_size+var_size), logits of one batch
+        """
+        node_stacks = [[TreeNode(init_hidden)] for init_hidden in problem_output.split(1, dim=0)]
+        embeddings_stacks = [[] for _ in range(encoder_outputs.size(0))]
+        left_child_trees = [None]*encoder_outputs.size(0)
+        all_node_outputs = []
+
+        for t in range(text_target.size(1)):
+            num_score, current_embeddings, current_context, current_all_embeddings = self.predict(
+                    node_stacks, 
+                    left_child_trees, 
+                    encoder_outputs, 
+                    self.embedding_var, 
+                    self.source_mask, 
+                    self.candi_mask,
+                    self.embedding_op_const)
+
+            all_node_outputs.append(num_score) # [B x (op_size+const_size+var_size)] * S
+
+            left_child, right_child, token_embedding = self.generate(
+                    current_embeddings, 
+                    text_target[:,t], 
+                    current_context, 
+                    self.embedding_op_const)
+
+            left_child_trees = []
+
+            for idx, (l, r, node_stack, target_id, embeddings_stack) in enumerate(zip(left_child.split(1), right_child.split(1),
+                                                                        node_stacks, text_target[:,t].tolist(), embeddings_stacks)):
+                # Determines whether the tree traversal is complete                                                    
+                if len(node_stack) != 0:
+                    node_stack.pop()
+                else:
+                    left_child_trees.append(None)
+                    continue
+                if target_id < self.op_num:
+                    node_stack.append(TreeNode(r))
+                    node_stack.append(TreeNode(l, left_flag=True))
+                    # embeddings_stack, put e(y|P) of op in temporarily
+                    embeddings_stack.append(TreeEmbedding(token_embedding[idx].unsqueeze(0), False))
+                else:
+                    current_num = current_all_embeddings[idx, target_id].unsqueeze(0) # 1 x H
+                    # Reach the right leaf node and merge the tree representation from bottom up
+                    while len(embeddings_stack) > 0 and embeddings_stack[-1].terminal: 
+                        sub_stree = embeddings_stack.pop()
+                        op = embeddings_stack.pop()
+                        # embedding vector of two sub-targets is merged as the subtree embedding of nodes, corresponding to eq(12)
+                        # with e(y|P), sub_tree_1 and sub_tree_2
+                        current_num = self.merge(op.embedding, sub_stree.embedding, current_num)
+                    embeddings_stack.append(TreeEmbedding(current_num, True))
+                # Reach the left leaf node and save the representation of the left subtree for generation of q
+                if len(embeddings_stack) > 0 and embeddings_stack[-1].terminal: 
+                    left_child_trees.append(embeddings_stack[-1].embedding)
+                else:
+                    left_child_trees.append(None)
+
+        all_node_outputs = torch.stack(all_node_outputs, dim=1)  
+
+        return all_node_outputs
+
+    def _forward_test(self, encoder_outputs, problem_output):
+
+        exp_outputs = [] 
+
+        for sample_id in range(encoder_outputs.size(0)):
+            # set batch size as 1
+            node_stacks = [[TreeNode(problem_output[sample_id:sample_id+1])]]
+            embeddings_stacks = [[]]
+            left_child_trees = [None]
+            beams = [TreeBeam(0.0, node_stacks, embeddings_stacks, left_child_trees, [])]
+
+            for _ in range(self.cfg.max_output_len):
+                # re-maintain of one beams
+                current_beams = []
+
+                while len(beams) > 0:
+                    beam_item = beams.pop()
+                    # The candidates are stored in beams in all process
+                    if len(beam_item.node_stacks[0]) == 0:
+                        current_beams.append(beam_item) 
+                        continue
+                    num_score, current_embeddings, current_context, current_all_embeddings = self.predict(
+                            beam_item.node_stacks, 
+                            beam_item.left_child_trees, 
+                            encoder_outputs[sample_id:sample_id+1], 
+                            self.embedding_var[sample_id:sample_id+1], 
+                            self.source_mask[sample_id:sample_id+1], 
+                            self.candi_mask[sample_id:sample_id+1],
+                            self.embedding_op_const)
+
+                    out_score = F.log_softmax(num_score, dim=1)
+                    topv, topi = out_score.topk(self.cfg.beam_size)
+
+                    for tv, ti in zip(topv.split(1, dim=1), topi.split(1, dim=1)):
+
+                        current_node_stack = copy_list(beam_item.node_stacks)
+                        current_left_child_trees = []
+                        current_embeddings_stacks = copy_list(beam_item.embeddings_stacks)
+                        current_out = copy.deepcopy(beam_item.out)
+
+                        out_token = int(ti)
+                        current_out.append(out_token)
+                        current_node_stack[0].pop()
+
+                        if out_token < self.op_num:
+                            generate_input = torch.LongTensor([out_token]).cuda()
+                            left_child, right_child, token_embedding = self.generate(
+                                current_embeddings, 
+                                generate_input, 
+                                current_context, 
+                                self.embedding_op_const)
+                            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(token_embedding, False))
+                        else:
+                            current_num = current_all_embeddings[:, out_token]
+                            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_child_trees.append(current_embeddings_stacks[0][-1].embedding)
+                        else:
+                            current_left_child_trees.append(None)
+                        
+                        current_beams.append(TreeBeam(beam_item.score+float(tv), current_node_stack, current_embeddings_stacks,
+                                                    current_left_child_trees, current_out))
+                
+                beams = sorted(current_beams, key=lambda x: x.score, reverse=True)
+                beams = beams[:self.cfg.beam_size]
+
+                # early termination
+                flag = True
+                for beam_item in beams:
+                    if len(beam_item.node_stacks[0]) != 0:
+                        flag = False
+                        break 
+                if flag: break
+
+            exp_outputs.append(beams[0].out)
+
+        return exp_outputs

model/encoder/__init__.py

@@ -0,0 +1,21 @@
+from .lstm import LSTM
+from .gru import GRU
+from config import encoder_list
+from .transformer import TransformerEncoder
+
+def get_encoder(params, *args):
+
+    if not params.encoder_type in encoder_list:
+        raise NotImplementedError(
+            "Unsupported Classifier: {}".format(params.encoder_type))
+
+    if params.encoder_type == "transformer":
+        pass
+    elif params.encoder_type == "lstm":
+        encoder = LSTM(params, *args)
+    elif params.encoder_type == "gru":
+        encoder = GRU(params, *args)
+    else:
+        raise NotImplementedError("Unsupported Encoder: {}".format(params.encoder_type))
+
+    return encoder

model/encoder/gru.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+
+
+class GRU(nn.Module):
+
+    def __init__(self, cfg):
+        super(GRU, self).__init__()
+
+        self.is_bidirectional = True
+        self.batch_first = True
+        self.gru = nn.GRU(
+            input_size = cfg.encoder_embedding_size,
+            hidden_size = cfg.encoder_hidden_size, # int(hidden_size / num_directions),
+            num_layers = cfg.encoder_layers,
+            bidirectional = self.is_bidirectional,
+            dropout = cfg.dropout_rate,
+            batch_first = self.batch_first
+        )
+        self.hidden_size = cfg.encoder_hidden_size
+        self.dropout = nn.Dropout(cfg.dropout_rate)
+    
+    def forward(self, src_emb, input_lengths, hidden=None):
+
+        input_emb = self.dropout(src_emb)
+        # input_emb = src_emb
+        packed = nn.utils.rnn.pack_padded_sequence(input_emb, input_lengths.cpu(), \
+                                            batch_first=self.batch_first, enforce_sorted=False)
+        pade_hidden = hidden
+        pade_outputs, pade_hidden = self.gru(packed, pade_hidden)
+        pade_outputs, _ = nn.utils.rnn.pad_packed_sequence(pade_outputs, batch_first=self.batch_first)
+        # pade_outputs [B, S, hidden_size*num_directions] 
+        # pade_hidden [n_layers*num_directions, B, hidden_size]
+        if self.is_bidirectional: 
+            pade_outputs = pade_outputs[:, :, :self.hidden_size] + pade_outputs[:, :, self.hidden_size:]  # B x S x H
+            pade_hidden = pade_hidden[0::2, :, :] + pade_hidden[1::2, :, :]
+
+        return pade_outputs, pade_hidden
+
+
+
+

model/encoder/lstm.py

@@ -0,0 +1,23 @@
+import torch.nn as nn
+
+
+class LSTM(nn.Module):
+    
+    def __init__(self, cfg):
+        super(LSTM, self).__init__()
+
+        self.lstm = nn.LSTM(
+            input_size=cfg.WORD_EMBED_SIZE,
+            hidden_size=cfg.HIDDEN_SIZE, # int(hidden_size / num_directions),
+            num_layers=cfg.NUM_LAYERS,
+            batch_first=cfg.BATCH_FIRST,  # first dim is batch_size or not
+            bidirectional=cfg.BIDIRECTIONAL
+        )
+
+    def forward(self, input, h0, c0):
+        output, (hn, cn) = self.lstm(input, (h0, c0))
+        return output, hn, cn
+
+
+
+

model/encoder/transformer.py

@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+from utils.utils import sequence_mask
+import math
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, max_len=5000, dropout=0.1):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        """
+            x: [B, max_len, d_model]
+            pe: [1, max_len, d_model]
+        """
+        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
+        return self.dropout(x)
+
+class LearnedPositionEncoding(nn.Module):
+
+    def __init__(self, d_model, max_len = 20):
+        super(LearnedPositionEncoding, self).__init__()
+        self.embedding = nn.Embedding(max_len, d_model)
+
+    def forward(self, x, var_pos):
+        """
+            x: [B, max_len, d_model]
+            var_pos: [B, var_len]
+        """
+        loc_mat = torch.zeros(x.size(0), x.size(1), dtype=torch.int64).cuda()
+        pos_id = torch.arange(1, var_pos.size(1)+1).repeat(var_pos.size(0), 1).cuda()
+        pos_id[var_pos==var_pos.min()] = 0
+        loc_mat.scatter_(1, var_pos, pos_id)
+
+        x = x + self.embedding(loc_mat)
+
+        return x
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, d_model=256, nhead=8, num_encoder_layers=6, dim_feedforward=1024, dropout=0.2):
+        super(TransformerEncoder,self).__init__()
+
+        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
+        encoder_norm = nn.LayerNorm(d_model)
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+        self.position = PositionalEncoding(d_model=d_model)
+        
+        self._reset_parameters()
+        self.d_model = d_model
+        self.nhead = nhead
+    
+    def _reset_parameters(self):
+        """
+            Initiate parameters in the transformer model.
+        """
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, len_src, emb_src):
+        # mask
+        src_key_padding_mask = ~sequence_mask(len_src)
+        # position encoding
+        emb_src = self.position(emb_src) 
+        # encoder   
+        memory = self.encoder(emb_src.permute(1,0,2), src_key_padding_mask=src_key_padding_mask)
+
+        return memory.permute(1,0,2)

model/module/__init__.py

@@ -0,0 +1,2 @@
+from .module_ops import *
+from .attention import *

model/module/attention.py

@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+
+class Score(nn.Module):
+    def __init__(self, input_size, hidden_size):
+        super(Score, self).__init__()
+        self.attn = nn.Linear(hidden_size + input_size, hidden_size)
+        self.score = nn.Linear(hidden_size, 1, bias=False)
+
+    def forward(self, hidden, candi_embeddings, candi_mask=None):
+        '''
+        Arguments:
+            hidden: B x 1 x 2H
+            candi_embeddings: B x candi_size x H
+            candi_mask: B x candi_size
+        Return:
+            score: B x candi_size
+        '''
+        hidden = hidden.repeat(1, candi_embeddings.size(1), 1)  # B x candi_size x H
+        # For each position of encoder outputs
+        energy_in = torch.cat((hidden, candi_embeddings), 2)  # B x candi_size x 3H
+        score = self.score(torch.tanh(self.attn(energy_in))).squeeze(-1)  # B x candi_size
+        if candi_mask is not None:
+            score = score.masked_fill_(~candi_mask, -1e12)
+        return score
+
+class Attn(nn.Module):
+    def __init__(self, input_size, hidden_size):
+        super(Attn, self).__init__()
+        self.attn = nn.Linear(hidden_size + input_size, hidden_size)
+        self.score = nn.Linear(hidden_size, 1, bias=False)
+
+    def forward(self, hidden, encoder_outputs, seq_mask=None):
+        '''
+        Arguments:
+            hidden: B x 1 x H (q)
+            encoder_outputs: B x S x H
+            seq_mask: B x S
+        Return:
+            attn_energies: B x S
+        '''
+        hidden = hidden.repeat(1, encoder_outputs.size(1), 1)  # B x S x H
+        energy_in = torch.cat((hidden, encoder_outputs), 2) # B x S x 2H
+        score_feature = torch.tanh(self.attn(energy_in)) # B x S x H
+        attn_energies = self.score(score_feature).squeeze(-1)  # B x S
+        if seq_mask is not None:
+            attn_energies = attn_energies.masked_fill_(~seq_mask, -1e12)
+        attn_energies = nn.functional.softmax(attn_energies, dim=1)  # B x S
+
+        return attn_energies
+    
+class Score_Multi(nn.Module):
+    def __init__(self, input_size, hidden_size):
+        super(Score_Multi, self).__init__()
+        self.attn = nn.Linear(hidden_size + input_size, hidden_size)
+        self.score = nn.Linear(hidden_size, 1, bias=False)
+
+    def forward(self, hidden, candi_embeddings, candi_mask=None):
+        '''
+        Arguments:
+            hidden: B x S x H
+            candi_embeddings: B x candi_size x H
+            candi_mask: B x candi_size
+        Return:
+            score: B x S x candi_size
+        '''
+        hidden = hidden.unsqueeze(2).repeat(1, 1, candi_embeddings.size(1), 1) # B x S x candi_size x H
+        candi_embeddings = candi_embeddings.unsqueeze(1).repeat(1, hidden.size(1), 1, 1) # B x S x candi_size x H
+        candi_mask = candi_mask.unsqueeze(1).repeat(1, hidden.size(1), 1) # B x S x candi_size
+        energy_in = torch.cat((hidden, candi_embeddings), -1)  # B x S x candi_size x 2H
+        score = self.score(torch.tanh(self.attn(energy_in))).squeeze(-1)  # B x S x candi_size
+        if candi_mask is not None:
+            score = score.masked_fill_(~candi_mask, -1e12)
+        return score

model/module/module_ops.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+
+
+class GAP(nn.Module):
+    """
+        Global Average pooling
+        Widely used in ResNet, Inception, DenseNet, etc.
+     """
+    def __init__(self):
+        super(GAP, self).__init__()
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+    def forward(self, x):
+        x = self.avgpool(x)
+        # x = x.view(x.shape[0], -1)
+        return x
+
+
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+

requirements.txt

@@ -0,0 +1,8 @@
+torch==1.7.1
+torchvision==0.8.2
+gradio==4.16.0
+Pillow>=9.0.0
+numpy>=1.19.0
+antlr4-python3-runtime==4.10
+sympy==1.11.1
+func_timeout==4.3.5

utils/__init__.py

@@ -0,0 +1,4 @@
+from .lr_scheduler import *
+from .utils import *
+
+

utils/lr_scheduler.py

@@ -0,0 +1,47 @@
+import torch
+from bisect import bisect_right
+
+
+class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(
+            self,
+            optimizer,
+            milestones,
+            gamma=0.1,
+            warmup_factor=1.0 / 3,
+            warmup_epochs=5,
+            warmup_method="linear",
+            last_epoch=-1,
+    ):
+        if not list(milestones) == sorted(milestones):
+            raise ValueError(
+                "Milestones should be a list of" " increasing integers. Got {}",
+                milestones,
+            )
+
+        if warmup_method not in ("constant", "linear"):
+            raise ValueError(
+                "Only 'constant' or 'linear' warmup_method accepted"
+                "got {}".format(warmup_method)
+            )
+        self.milestones = milestones
+        self.gamma = gamma
+        self.warmup_factor = warmup_factor
+        self.warmup_epochs = warmup_epochs
+        self.warmup_method = warmup_method
+        super(WarmupMultiStepLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        warmup_factor = 1
+        if self.last_epoch < self.warmup_epochs:
+            if self.warmup_method == "constant":
+                warmup_factor = self.warmup_factor
+            elif self.warmup_method == "linear":
+                alpha = float(self.last_epoch) / self.warmup_epochs
+                warmup_factor = self.warmup_factor * (1 - alpha) + alpha
+        return [
+            base_lr
+            * warmup_factor
+            * self.gamma ** bisect_right(self.milestones, self.last_epoch)
+            for base_lr in self.base_lrs
+        ]

utils/utils.py

@@ -0,0 +1,369 @@
+import os
+import torch
+from utils.lr_scheduler import WarmupMultiStepLR
+from config import *
+import datetime
+import torch.distributed as dist
+from datasets.operators import result_compute, normalize_exp
+from func_timeout import func_timeout
+import random
+import gc
+
+def save_checkpoint(state, is_best, dump_path=None):
+    if is_best:
+        dump_path_best = os.path.join(dump_path, 'best_model.pth')
+        torch.save(state, dump_path_best)
+    else:
+        dump_path_recent = os.path.join(dump_path, str(state['epoch'])+'.pth')
+        torch.save(state, dump_path_recent)
+
+class AverageMeter(object):
+    """
+    Computes and stores the average and current value
+    """
+    def __init__(self, name, fmt=':f'):
+        self.name = name
+        self.fmt = fmt
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
+        return fmtstr.format(**self.__dict__)
+
+class ProgressMeter(object):
+    def __init__(self, num_batches, meters, args, prefix=""):
+        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
+        self.meters = meters
+        self.prefix = prefix
+        self.args = args
+
+    def display(self, batch, lr=None):
+        entries = [self.prefix + self.batch_fmtstr.format(batch)]
+        entries += [str(meter) for meter in self.meters]
+        if not lr is None:
+            entries += ["lr: "+str(format(lr, '.6f'))]
+        self.args.logger.info('\t'.join(entries))
+ 
+    def _get_batch_fmtstr(self, num_batches):
+        num_digits = len(str(num_batches // 1))
+        fmt = '{:' + str(num_digits) + 'd}'
+        return '[' + fmt + '/' + fmt.format(num_batches) + ']'
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """
+        Sets the learning rate to the initial LR decayed by 10 every 30 epochs
+    """
+    lr = args.lr * (0.1**(epoch // 30))
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+def accuracy(output, target, topk=(1, )):
+    """
+        Computes the accuracy over the k top predictions for the specified values of k
+    """
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].contiguous().view(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+def get_scheduler(args, optimizer):
+    if args.scheduler_type == "multistep":
+        scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            args.scheduler_step,
+            gamma=args.scheduler_factor,
+        )
+    elif args.scheduler_type == "cosine":
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=args.max_epochs, eta_min=1e-6)
+    elif args.scheduler_type == "warmup":
+        scheduler = WarmupMultiStepLR(
+            optimizer,
+            args.scheduler_step,
+            gamma=args.scheduler_factor,
+            warmup_epochs=args.warm_epoch,
+        )
+    else:
+        raise NotImplementedError("Unsupported LR Scheduler: {}".format(args.scheduler_type))
+    return scheduler
+
+def get_optimizer(args, model):
+    if args.use_MLM_pretrain:
+        pretrain_params = list(map(id, model.mlm_pretrain.parameters()))
+        other_params = filter(lambda p: id(p) not in pretrain_params, model.parameters())
+
+    if args.optimizer_type == "SGD":
+        optimizer = torch.optim.SGD(
+            model.parameters(),
+            lr=args.lr,
+            momentum=args.momentum,
+            weight_decay=args.weight_decay,
+            nesterov=True,
+        )
+    elif args.optimizer_type == "ADAM":
+        if args.use_MLM_pretrain:
+            optimizer = torch.optim.Adam(
+                [{"params":model.mlm_pretrain.parameters()},
+                    {"params":other_params, "lr":args.lr_LM}],
+                lr=args.lr,
+                betas=(0.9, 0.999),
+                weight_decay=args.weight_decay,
+            )
+        else:
+            optimizer = torch.optim.Adam(
+                model.parameters(),
+                lr=args.lr,
+                betas=(0.9, 0.999),
+                weight_decay=args.weight_decay,
+            )
+    elif args.optimizer_type == "ADAMW":
+        if args.use_MLM_pretrain:
+            optimizer = torch.optim.AdamW(
+                [{"params":model.mlm_pretrain.parameters(), "lr":args.lr_LM},
+                    {"params":other_params}],
+                lr=args.lr,
+                weight_decay=args.weight_decay,
+            )
+        else:
+            optimizer = torch.optim.AdamW(
+                model.parameters(),
+                lr=args.lr,
+                weight_decay=args.weight_decay,
+            )
+    else:
+        raise NotImplementedError("Unsupported Optimizer Type : {}".format(args.optimizer_type)) 
+
+    return optimizer
+
+def reduce_mean(tensor, nprocs):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= nprocs
+    return rt
+
+def set_cuda(data_dict):
+    for key in data_dict:
+        if torch.is_tensor(data_dict[key]):
+            data_dict[key] = data_dict[key].cuda()
+
+def initialize_logger(params, ):
+    """
+        Initialize the experience:
+        - dump parameters
+        - create a logger
+    """
+    while True:
+        exp_id = datetime.datetime.strftime(datetime.datetime.now(),'%Y-%m-%d-%H-%M-%S')
+        if not os.path.exists(os.path.join(params.dump_path, exp_id)):
+            break
+    params.dump_path = os.path.join(params.dump_path, exp_id)
+    if params.local_rank == 0:
+        os.makedirs(params.dump_path)
+    # create a logger
+    logger = create_logger(os.path.join(params.dump_path,'record.log'), params.local_rank)
+    logger.info("============ Initialized logger ============")
+    logger.info("\n"+"\n".join("\t\t\t\t%s: %s" % (k, str(v))
+                          for k, v in sorted(dict(vars(params)).items())))
+    logger.info("The experiment results will be stored in %s" % params.dump_path)
+    return logger
+
+def aeq(*args):
+    """
+    Assert all arguments have the same value
+    """
+    arguments = (arg for arg in args)
+    first = next(arguments)
+    assert all(arg == first for arg in arguments), \
+        "Not all arguments have the same value: " + str(args)
+
+def sequence_mask(lengths, max_len=None):
+    """
+    Creates a boolean mask from sequence lengths.
+    """
+    batch_size = lengths.numel()
+    max_len = max_len or lengths.max()
+    return torch.arange(0, max_len, device=lengths.device) \
+            .type_as(lengths) \
+            .repeat(batch_size, 1) \
+            .lt(lengths.unsqueeze(1))
+
+def copy_list(l):
+    r = []
+    if len(l) == 0:
+        return r
+    for i in l:
+        if type(i) is list:
+            r.append(copy_list(i))
+        else:
+            r.append(i)
+    return r
+
+def compute_exp_result_choice(test_preds, var_dict, exp_dict, tgt_lang):
+    
+    """
+    Arguments
+        test_preds: B x candi_size(beam_size) x token_list
+        var_dict: {'pos', 'len', 'var_value', 'arg_value'}
+        exp_dict: {'exp', 'len', 'answer'}
+        tgt_lang: vocab of target text
+    Returns:
+        ans_acc 
+        eq_acc
+    """
+    gc.collect()
+    ans_num = eq_num = 0
+
+    for k in range(len(test_preds)): # batch id 
+        tgt = exp_dict['exp'][k][1:exp_dict['len'][k]-1].tolist() # Remove special symbols [SOS] and [EOS]
+        var2arg_dict = {'N'+str(i+len(var_dict['var_value'][k])):item \
+                                for i, item in enumerate(var_dict['arg_value'][k])}
+        tgt = tgt_lang.sentence_from_indexes(tgt, var2arg_dict)
+        num_list = var_dict['var_value'][k]
+        tgt_result = float(exp_dict['answer'][k])
+        choices = exp_dict['choices'][k]
+        is_find_ans = False
+
+        for j in range(len(test_preds[k])): # pred candi id 
+            try:
+                pred = tgt_lang.sentence_from_indexes(test_preds[k][j], var2arg_dict)
+                pred = normalize_exp(pred)
+                pred_result = float(func_timeout(2.0, result_compute, \
+                                        kwargs=dict(num_all_list=num_list, exp_tokens=pred)))
+                if pred == tgt:
+                    ans_num += 1
+                    eq_num += 1
+                    is_find_ans = True
+                    break
+                for item in choices:
+                    if abs(pred_result-item)<5e-2: 
+                        is_find_ans = True
+                if  is_find_ans and abs(pred_result-tgt_result)<5e-3:
+                    ans_num +=1
+                    if len(pred)==len(tgt):
+                        eq_num += 1
+                if  is_find_ans: break
+            except:
+                pass
+        
+        if not is_find_ans:
+            pred_result = random.choice(choices)
+            if  abs(pred_result-tgt_result)<5e-2:
+                ans_num +=1
+                
+    return ans_num/len(test_preds), eq_num/len(test_preds)
+
+def compute_exp_result_topk(test_preds, var_dict, exp_dict, tgt_lang, k_num = 3):
+    
+    """
+    Arguments
+        test_preds: B x candi_size(beam_size) x token_list
+        var_dict: {'pos', 'len', 'var_value', 'arg_value'}
+        exp_dict: {'exp', 'len', 'answer'}
+        tgt_lang: vocab of target text
+    Returns:
+        ans_acc 
+        eq_acc
+    """
+    gc.collect()
+    ans_num = eq_num = 0
+
+    for k in range(len(test_preds)): # batch id 
+        tgt = exp_dict['exp'][k][1:exp_dict['len'][k]-1].tolist() # Remove special symbols [SOS] and [EOS]
+        var2arg_dict = {'N'+str(i+len(var_dict['var_value'][k])):item \
+                                for i, item in enumerate(var_dict['arg_value'][k])}
+        tgt = tgt_lang.sentence_from_indexes(tgt, var2arg_dict)
+        num_list = var_dict['var_value'][k]
+        tgt_result = float(exp_dict['answer'][k])
+        is_ans_same = is_eq_same = False
+
+        for j in range(k_num): # top-n
+            try:
+                pred = tgt_lang.sentence_from_indexes(test_preds[k][j], var2arg_dict)
+                pred = normalize_exp(pred)
+                pred_result = float(func_timeout(2.0, result_compute, \
+                                        kwargs=dict(num_all_list=num_list, exp_tokens=pred)))
+                if pred == tgt:
+                    is_ans_same = True
+                    is_eq_same = True
+                    break
+                if abs(pred_result-tgt_result)<5e-3: 
+                    is_ans_same = True
+                    if len(pred)==len(tgt):
+                        is_eq_same = True
+                        break
+            except:
+                pass
+        
+        if is_ans_same: ans_num +=1
+        if is_eq_same: eq_num +=1
+        
+    return ans_num/len(test_preds), eq_num/len(test_preds)
+
+
+def compute_exp_result_comp(test_preds, var_dict, exp_dict, tgt_lang):
+    
+    """
+    Arguments
+        test_preds: B x candi_size(beam_size) x token_list
+        var_dict: {'pos', 'len', 'var_value', 'arg_value'}
+        exp_dict: {'exp', 'len', 'answer'}
+        tgt_lang: vocab of target text
+    Returns:
+        ans_acc 
+        eq_acc
+    """
+    gc.collect()
+    ans_num = eq_num = 0
+
+    for k in range(len(test_preds)): # batch id 
+        tgt = exp_dict['exp'][k][1:exp_dict['len'][k]-1].tolist() # Remove special symbols [SOS] and [EOS]
+        var2arg_dict = {'N'+str(i+len(var_dict['var_value'][k])):item \
+                                for i, item in enumerate(var_dict['arg_value'][k])}
+        tgt = tgt_lang.sentence_from_indexes(tgt, var2arg_dict)
+        num_list = var_dict['var_value'][k]
+        tgt_result = float(exp_dict['answer'][k])
+        is_ans_same = is_eq_same = False
+
+        for j in range(len(test_preds[k])): # pred candi id 
+            try:
+                pred = tgt_lang.sentence_from_indexes(test_preds[k][j], var2arg_dict)
+                pred = normalize_exp(pred)
+                pred_result = float(func_timeout(2.0, result_compute, \
+                        kwargs=dict(num_all_list=num_list, exp_tokens=pred)))
+                if pred == tgt:
+                    is_ans_same = True
+                    is_eq_same = True
+                    break
+                if abs(pred_result-tgt_result)<5e-3:
+                    is_ans_same = True
+                    if len(pred)==len(tgt):
+                        is_eq_same = True
+                break
+            except:
+                pass
+        
+        if is_ans_same: ans_num +=1
+        if is_eq_same: eq_num +=1
+        
+    return ans_num/len(test_preds), eq_num/len(test_preds)

vocab/vocab_src.txt

@@ -0,0 +1,322 @@
+[PAD]
+[UNK]
+[CLS]
+[MASK]
++
+,
+.
+/
+1
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+2
+20
+3
+4
+5
+6
+7
+8
+9
+=
+?
+A
+B
+B'
+B0
+B1
+C
+C0
+C1
+D
+E
+F
+F0
+F1
+G
+G0
+G1
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+Q'
+R
+S
+T
+U
+V
+W
+W'
+X
+Y
+Z
+\angle
+\cong
+\cos
+\odot
+\parallel
+\parallelogram
+\perp
+\phi
+\sim
+\sin
+\tan
+\triangle
+\widehat
+a
+about
+all
+altitude
+altitudes
+an
+and
+angle
+angles
+any
+appear
+appears
+are
+area
+areas
+as
+assume
+at
+b
+base
+bases
+be
+below
+between
+bisector
+bisectors
+bisects
+blue
+both
+by
+c
+calculator
+center
+centers
+centimeters
+central
+centroid
+chord
+chords
+circle
+circles
+circumference
+circumscribed
+circumscribes
+cm
+cm^{2}
+collinear
+common
+complementary
+composite
+congruent
+connecting
+corners
+cosines
+cut
+d
+degree
+determine
+diagonal
+diagonals
+diagram
+diameter
+diameters
+distance
+drawn
+e
+each
+elm
+equal
+equidistant
+equilateral
+exact
+express
+f
+factor
+feet
+figure
+figures
+find
+for
+form
+formed
+four
+from
+ft
+ft^{2}
+g
+given
+green
+h
+half
+has
+have
+having
+height
+hexagon
+how
+hypotenuse
+i
+if
+in
+in^{2}
+incenter
+inches
+inscribed
+inside
+intersect
+intersected
+intersecting
+into
+is
+isosceles
+it
+its
+j
+k
+kite
+l
+law
+legs
+length
+lengths
+let
+lieson
+line
+linear
+lines
+long
+m
+m^{2}
+major
+make
+makes
+measure
+measurement
+measures
+median
+medians
+meet
+meters
+midpoint
+midpoints
+midsegment
+midsegments
+miles
+millimeters
+minor
+mm
+mm^{2}
+must
+n
+o
+octagon
+of
+off
+on
+one
+otherwise
+p
+pair
+parallel
+parallelogram
+pentagon
+pentagons
+perimeter
+perimeters
+perpendicular
+plum
+point
+points
+polygon
+polygons
+proportion
+pythagorean
+q
+quadrilateral
+r
+radius
+ratio
+ray
+rectangle
+red
+refer
+region
+regular
+respectively
+rhombus
+right
+s
+scale
+sector
+segment
+segments
+shaded
+shown
+side
+sides
+similar
+sines
+so
+solve
+special
+square
+stated
+straight
+such
+sum
+suppose
+t
+tangent
+tangents
+that
+the
+theorem
+this
+times
+to
+trapezoid
+triangle
+triangles
+triple
+twice
+two
+u
+units
+unless
+use
+v
+value
+variable
+vertex
+w
+what
+where
+which
+with
+would
+x
+y
+yards
+yd^{2}
+z

vocab/vocab_tgt.txt

@@ -0,0 +1,67 @@
+[PAD]
+[SOS]
+[EOS]
+V0
+V1
+V2
+V3
+V4
+V5
+V6
+C0.5
+C2
+C3
+C4
+C5
+C6
+C8
+C60
+C90
+C180
+C360
+ArcSeg_Area
+Chord2_Ang
+Circle_D_Area
+Circle_D_Circum
+Circle_R_Area
+Circle_R_Circum
+Cos_Law
+Equal
+Gcos
+Geo_Mean
+Get
+Gougu
+Gsin
+Gtan
+Iso_Tri_Ang
+Kite_Area
+Median
+Multiple
+Ngon_Angsum
+PRK_Perim
+Para_Area
+Proportion
+RNgon_B_Area
+RNgon_H_Area
+RNgon_L_Area
+Ratio
+Rect_Area
+Rhom_Area
+Sin_Law
+Sum
+TanSec_Ang
+Trap_Area
+Tria_BH_Area
+Tria_SAS_Area
+N0
+N1
+N2
+N3
+N4
+N5
+N6
+N7
+N8
+N9
+N10
+N11