def __init__(self, class_mapping, dictionary=None, mode='all', bert_weight=None):
    assert (mode in ['all', 'bert', 'trie', 'ed']), 'Mode is not supported'
    self.mode = mode
    self.dictionary = dictionary
    self.class_mapping = class_mapping
    self.idx_mapping = {v: k for (k, v) in class_mapping.items()}
    if (self.mode == 'bert'):
        self.use_bert = True
    if (self.mode == 'trie'):
        self.use_trie = True
    if (self.mode == 'ed'):
        self.use_ed = True
    if (self.mode == 'all'):
        self.use_bert = True
        self.use_trie = True
        self.use_ed = True
    if self.use_bert:
        self.bert = retrieval.PhoBERT(self.idx_mapping, bert_weight)
    if self.use_ed:
        self.ed = retrieval.get_heuristic_retrieval('diff')
    if self.use_trie:
        self.trie = retrieval.get_heuristic_retrieval('trie')
    if (self.use_ed or self.use_trie):
        if (self.dictionary is None):
            self.dictionary = {}
            df = pd.read_csv('./modules/retrieval/heuristic/custom-dictionary.csv')
            for (id, row) in df.iterrows():
                self.dictionary[row.text.lower()] = row.lbl
def __init__(self, class_mapping, dictionary=None, mode='all', bert_weight=None):
    assert (mode in ['all', 'bert', 'trie', 'ed']), 'Mode is not supported'
    self.mode = mode
    self.dictionary = dictionary
    self.class_mapping = class_mapping
    self.idx_mapping = {v: k for (k, v) in class_mapping.items()}
    if (self.mode == 'bert'):
        self.use_bert = True
    if (self.mode == 'trie'):
        self.use_trie = True
    if (self.mode == 'ed'):
        self.use_ed = True
    if (self.mode == 'all'):
        self.use_bert = True
        self.use_trie = True
        self.use_ed = True
    if self.use_bert:
        self.bert = retrieval.PhoBERT(self.idx_mapping, bert_weight)
    if self.use_ed:
        self.ed = retrieval.get_heuristic_retrieval('diff')
    if self.use_trie:
        self.trie = retrieval.get_heuristic_retrieval('trie')
    if (self.use_ed or self.use_trie):
        if (self.dictionary is None):
            self.dictionary = {}
            df = pd.read_csv('./modules/retrieval/heuristic/custom-dictionary.csv')
            for (id, row) in df.iterrows():
                self.dictionary[row.text.lower()] = row.lbl