def prompt_eval():
    eval_prompt = """You are an expert professor specialized in grading whether the prediction to the question is correct or not according to the real answer.
    ==================
    For example:
    ==================
    Question: What company owns the property of Marvel Comics?
    Answer: The Walt Disney Company
    Prediction: The Walt Disney Company
    Return: 1
    ==================
    Question: Which constituent college of the University of Oxford endows four professorial fellowships for sciences including chemistry and pure mathematics?
    Answer: Magdalen College
    Prediction: Magdalen College.
    Return: 1
    ==================
    Question: Which year was Marvel started?
    Answer: 1939
    Prediction: 1200
    Return: 0
    ==================
    You are grading the following question:
    Question: {question}
    Answer: {answer}
    Prediction: {prediction}
    If the prediction is correct according to answer, return 1. Otherwise, return 0.
    Return: your reply can only be one number '0' or '1'
    """

    return eval_prompt

def knn_graph(i_d, k_knn, embs, strategy='cos'):
    idx, d = i_d

    emb = embs[idx]

    # build a knn Graph
    if strategy == 'cos':
        sim = cosine_similarity(emb, emb)

    elif strategy == 'dp':
        sim = np.matmul(emb, emb.transpose(1, 0))

    # topk
    top_idx = np.argsort(-sim, axis=1)[:, 1:k_knn + 1]

    tail_nodes = np.arange(top_idx.shape[0]).repeat(k_knn) # flatten
    head_nodes = top_idx.reshape(-1)
    edges = [(node1, node2) for node1, node2 in zip(tail_nodes, head_nodes)]

    G = nx.DiGraph()
    G.add_edges_from(edges)

    return idx, G
上述代碼實現了，兩個節點根據它倆之間向量相似度構建邊。

def mp_loss(model, batch):
    embs = model(batch)
    loss_fct = CrossEntropyLoss(ignore_index = -1)

    c_embs = torch.cat([embs["c1_emb"], embs["c2_emb"]], dim = 0) # 2B x d
    n_embs = torch.cat([embs["n1_emb"].unsqueeze(1), embs["n2_emb"].unsqueeze(1)], dim = 1) # B*2*M*h

    scores_1 = torch.mm(embs["q_emb"], c_embs.t()) # B x 2B
    n_scores_1 = torch.bmm(embs["q_emb"].unsqueeze(1), n_embs.permute(0, 2, 1)).squeeze(1) # B x 2B
    scores_2 = torch.mm(embs["q_c1_emb"], c_embs.t()) # B x 2B
    n_scores_2 = torch.bmm(embs["q_c1_emb"].unsqueeze(1), n_embs.permute(0, 2, 1)).squeeze(1) # B x 2B

    # mask the 1st hop
    bsize = embs["q_emb"].size(0)
    scores_1_mask = torch.cat([torch.zeros(bsize, bsize), torch.eye(bsize)], dim=1).to(embs["q_emb"].device)
    scores_1 = scores_1.float().masked_fill(scores_1_mask.bool(), float('-inf')).type_as(scores_1)
    scores_1 = torch.cat([scores_1, n_scores_1], dim=1)
    scores_2 = torch.cat([scores_2, n_scores_2], dim=1)

    target_1 = torch.arange(embs["q_emb"].size(0)).to(embs["q_emb"].device)
    target_2 = torch.arange(embs["q_emb"].size(0)).to(embs["q_emb"].device) + embs["q_emb"].size(0)

    loss = loss_fct(scores_1, target_1) + loss_fct(scores_2, target_2)

    return loss

class KG_retriever(object):
    def __init__(self, k):
        self.k = k

    def retrieve(self, data, G):
        corpus = [c for _, c in data['title_chunks']]
        candidates_idx = list(range(len(corpus)))

        seed = data['question']
        retrieve_idxs = []

        prev_length = 0
        count = 0
        retrieve_num = [10, 5, 5, 5, 3, 2, 2, 2, 2, 2, 2]
        while len(retrieve_idxs) < self.k:
            idxs = tf_idf(seed, candidates_idx, corpus, k = retrieve_num[count], visited = retrieve_idxs)
            retrieve_idxs.extend(idxs[:max(0, self.k - len(retrieve_idxs))])

            candidates_idx = set(chain(*[list(G.neighbors(node)) for node in idxs]))
            candidates_idx = list(candidates_idx.difference(retrieve_idxs))

            if len(retrieve_idxs) == prev_length:
                break
            else:
                prev_length = len(retrieve_idxs)

            count += 1

        return [corpus[idx] for idx in retrieve_idxs], None, None, None