Example #1
0
 def predict(self, mini_batch, verbose=False):
     kg, pn = self.kg, self.mdl
     e1, e2, r = self.format_batch(mini_batch)
     beam_search_output = search.beam_search(pn, e1, r, e2, kg,
                                             self.num_rollout_steps,
                                             self.beam_size)
     pred_e2s = beam_search_output['pred_e2s']
     pred_e2_scores = beam_search_output['pred_e2_scores']
     if verbose:
         # print inference paths
         search_traces = beam_search_output['search_traces']
         output_beam_size = min(self.beam_size, pred_e2_scores.shape[1])
         for i in range(len(e1)):
             for j in range(output_beam_size):
                 ind = i * output_beam_size + j
                 if pred_e2s[i][j] == kg.dummy_e:
                     break
                 search_trace = []
                 for k in range(len(search_traces)):
                     search_trace.append((int(search_traces[k][0][ind]),
                                          int(search_traces[k][1][ind])))
                     print('beam {}: score = {} \n<PATH> {}'.format(
                         j, float(pred_e2_scores[i][j]),
                         ops.format_path(search_trace, kg)))
     with torch.no_grad():
         pred_scores = zeros_var_cuda([len(e1), kg.num_entities])
         for i in range(len(e1)):
             pred_scores[i][pred_e2s[i].long()] = torch.exp(
                 pred_e2_scores[i])
     return pred_scores, beam_search_output["rule_score"]
Example #2
0
    def predict(self, mini_batch, verbose=False):
        kg, pn = self.kg, self.mdl
        e1, e2, r = self.format_batch(mini_batch)
        width = kg.num_entities - 1

        beam_search_output = search.beam_search(pn, e1, r, e2, kg,
                                                self.num_rollout_steps,
                                                self.beam_size)
        with torch.no_grad():
            pred_e2s = beam_search_output['pred_e2s']
            pred_e2_scores = beam_search_output['pred_e2_scores']
            pred_traces = beam_search_output['pred_traces']
            for u in range(pred_e2s.size()[0]):
                for v in range(pred_e2s.size()[1]):
                    if pred_e2s[u][v].item() == kg.dummy_end_e:
                        for i in range(self.num_rollout_steps, 0, -1):
                            if pred_traces[i][0][1][u * width +
                                                    v] != kg.dummy_end_e:
                                pred_e2s[u][v] = pred_traces[i][0][1][u * width
                                                                      + v]
                                break

        accum_scores = ops.sort_by_path(pred_traces, kg)

        # for testing
        ggg = open("traces.txt", "a")
        hhh = open("accum_scores.txt", "a")
        for i in range(len(e1)):
            if e1[i] == 104 and r[i] == 5 and e2[i] == 72:
                print("Epoch: ",
                      pred_traces[1][0][0][i * width:(i + 1) * width],
                      "\n",
                      pred_traces[1][0][1][i * width:(i + 1) * width],
                      "\n",
                      pred_traces[2][0][0][i * width:(i + 1) * width],
                      "\n",
                      pred_traces[2][0][1][i * width:(i + 1) * width],
                      file=ggg)
                print(sorted(accum_scores[i].items(),
                             key=lambda x: x[1],
                             reverse=True),
                      file=hhh)
        ggg.close()
        hhh.close()

        if verbose:
            # print inference paths
            search_traces = beam_search_output['search_traces']
            output_beam_size = min(self.beam_size, pred_e2_scores.shape[1])
            for i in range(len(e1)):
                for j in range(output_beam_size):
                    ind = i * output_beam_size + j
                    if pred_e2s[i][j] == kg.dummy_e:
                        break
                    search_trace = []
                    for k in range(len(search_traces)):
                        search_trace.append((int(search_traces[k][0][ind]),
                                             int(search_traces[k][1][ind])))
                    print('beam {}: score = {} \n<PATH> {}'.format(
                        j, float(pred_e2_scores[i][j]),
                        ops.format_path(search_trace, kg)))
        with torch.no_grad():
            uu = open("scores_pred.txt", "a")
            pred_scores = zeros_var_cuda([len(e1), kg.num_entities])
            for i in range(len(e1)):
                pred_scores[i][pred_e2s[i]] = torch.exp(pred_e2_scores[i])
                if e1[i] == 104 and r[i] == 5 and e2[i] == 72:
                    print(pred_scores[i], file=uu)

                # 计算ranking的方法有待检讨。理论上应该使用第二种
                # 这是两个加和取平均的方式。但实际上只有rma发挥了有效作用
                # pred_scores[i][pred_e2s_1[i]] += torch.div(pred_e2_scores_1[i], 2.0)
                # pred_scores[i][pred_e2s_2[i]] += torch.div(pred_e2_scores_2[i], 2.0)
                # pred_scores[i] = torch.exp(pred_scores[i])

                # 这是两步计算平均的方式
                # for j in range(len(pred_e2s_2)):
                #    pred_scores[i][pred_e2s_2[i][j]] = torch.exp(pred_e2_scores_2[i][j]) *\
                #                                ops.check_path(accum_scores[i],pred_traces_2, i*width+j)
        uu.close()
        return pred_scores