Fix Lang classes, CUDA compatibility, and config imports

app.py

@@ -13,43 +13,11 @@ 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.preprossing import SN, SrcLang, TgtLang
 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):
@@ -89,18 +57,9 @@ class Config:
 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())
+    # Load vocabularies using proper Lang classes
+    src_lang = SrcLang('./vocab/vocab_src.txt')
+    tgt_lang = TgtLang('./vocab/vocab_tgt.txt')
     
     # Create model
     model = Network(cfg, src_lang, tgt_lang)

loss/__init__.py

@@ -1,5 +1,6 @@
 from .loss import *
-from config import criterion_list
+
+criterion_list = ["CrossEntropy", "FocalLoss", "MaskedCrossEntropy"]
 
 
 def get_criterion(args):   

model/classifier/__init__.py

@@ -1,5 +1,6 @@
 from .classifier_ops import *
-from config import classifier_list
+
+classifier_list = ["FCNorm", "CosNorm", "DotProduct", "DistFC"]
 
        
 def get_classifier(args):

model/decoder/rnn_decoder.py

@@ -23,7 +23,8 @@ class DecoderRNN(nn.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.no_var_id = torch.arange(self.var_start).unsqueeze(0).cuda()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.no_var_id = torch.arange(self.var_start).unsqueeze(0).to(self.device)
         self.cfg = cfg
     
     def get_var_encoder_outputs(self, encoder_outputs, var_pos):
@@ -127,15 +128,15 @@ class DecoderRNN(nn.Module):
             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
+                    input_token = torch.LongTensor([[self.sos_id]]*rem_size).to(self.device) # 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_score = torch.FloatTensor([[0.0]]*rem_size).to(self.device) # rem_size x 1
                     current_exp_list = [[]]*rem_size
                 else:
-                    input_token = torch.LongTensor(token_list).unsqueeze(1).cuda() 
+                    input_token = torch.LongTensor(token_list).unsqueeze(1).to(self.device) 
                     rnn_hidden = rnn_hidden[:, cand_list]
                     rem_size = len(exp_list)
-                    current_score = torch.FloatTensor(score_list[:rem_size]).unsqueeze(1).cuda()
+                    current_score = torch.FloatTensor(score_list[:rem_size]).unsqueeze(1).to(self.device)
                     current_exp_list = exp_list
                     
                 # input embedding

utils/utils.py

@@ -1,7 +1,6 @@
 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