# Path: smore/common/config.py
# def parse_args(args=None):


import numpy as np
import random
import torch
import time
import os.path as osp
import torch.multiprocessing as mp
import traceback
import logging
import os
import shutil
import zipfile
import urllib.request as ur
from collections import defaultdict
from functools import wraps
from torch.multiprocessing import Queue
from _thread import start_new_thread
from tqdm import tqdm
from smore.common.config import name_query_dict, query_name_dict

                embedding = all_ent_embedding[:,ent_idx,:]
                offset_embedding = torch.zeros_like(embedding)
                idx += 1
                ent_idx += 1
            else:
                embedding, offset_embedding, idx, ent_idx = embed_query_new(queries, query_structure[0], idx, all_ent_embedding, ent_idx)
            for i in range(len(query_structure[-1])):
                r_embedding = main_relation_embedding[queries[:, idx]]
                r_offset_embedding = main_offset_embedding[queries[:, idx]]
                embedding += r_embedding
                offset_embedding += r_offset_embedding
                idx += 1
        else:
            embedding_list = []
            offset_embedding_list = []
            for i in range(len(query_structure)):
                embedding, offset_embedding, idx, ent_idx = embed_query_new(queries, query_structure[i], idx, all_ent_embedding, ent_idx)
                embedding_list.append(embedding)
                offset_embedding_list.append(offset_embedding)
            embedding = torch.mean(torch.stack(embedding_list), dim=0)[0]
            offset_embedding = torch.max(torch.stack(offset_embedding_list), dim=0)[0]

        return embedding, offset_embedding, idx, ent_idx


    def transform_union_query(queries, query_structure):
        '''
        transform 2u queries to two 1p queries
        transform up queries to two 2p queries
        '''

# Path: smore/common/config.py
# def parse_args(args=None):


# Path: smore/common/util.py
import numpy as np
import random
import torch
import time
import os.path as osp
import torch.multiprocessing as mp
import traceback
import logging
import os
import shutil
import zipfile
import urllib.request as ur
from collections import defaultdict
from functools import wraps
from torch.multiprocessing import Queue
from _thread import start_new_thread
from tqdm import tqdm
from smore.common.config import name_query_dict, query_name_dict

                embedding = all_ent_embedding[:,ent_idx,:]
                offset_embedding = torch.zeros_like(embedding)
                idx += 1
                ent_idx += 1
            else:
                embedding, offset_embedding, idx, ent_idx = embed_query_new(queries, query_structure[0], idx, all_ent_embedding, ent_idx)
            for i in range(len(query_structure[-1])):
                r_embedding = main_relation_embedding[queries[:, idx]]
                r_offset_embedding = main_offset_embedding[queries[:, idx]]
                embedding += r_embedding
                offset_embedding += r_offset_embedding
                idx += 1
        else:
            embedding_list = []
            offset_embedding_list = []
            for i in range(len(query_structure)):
                embedding, offset_embedding, idx, ent_idx = embed_query_new(queries, query_structure[i], idx, all_ent_embedding, ent_idx)
                embedding_list.append(embedding)
                offset_embedding_list.append(offset_embedding)
            embedding = torch.mean(torch.stack(embedding_list), dim=0)[0]
            offset_embedding = torch.max(torch.stack(offset_embedding_list), dim=0)[0]

        return embedding, offset_embedding, idx, ent_idx


    def transform_union_query(queries, query_structure):
        '''
        transform 2u queries to two 1p queries
        transform up queries to two 2p queries
        '''

        if query_name_dict[query_structure] == '2u-DNF':
