123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213 |
- ''' Translate input text with trained model. '''
-
- import torch
- import random
- import argparse
- from tqdm import tqdm
- import numpy as np
- from transformer.Translator_new import Translator
- import transformer.Constants as Constants
- from DataLoader_Feat import DataLoader
- import pickle
-
- idx2vec_addr = '/media/external_3TB/3TB/ramezani/pmoini/Trial/data_d8/idx2vec.pickle'
-
-
- def getF1(ground_truth, pred_cnt):
- right = np.dot(ground_truth, pred_cnt)
- pred = np.sum(pred_cnt)
- total = np.sum(ground_truth)
- print(right, pred, total)
-
- if pred == 0:
- return 0, 0, 0
- precision = right / pred
- recall = right / total
- if precision == 0 or recall == 0:
- return 0, 0, 0
- return (2 * precision * recall) / (precision + recall), precision, recall
-
-
- def getMSE(ground_truth, pred_cnt):
- return mean_squared_error(ground_truth, pred_cnt)
-
-
- def getMAP(ground_truth, pred_cnt):
- size_cascade = list(ground_truth).count(1)
- avgp = 0.
- ite_list = [[idx, item, ground_truth[idx]] for idx, item in enumerate(pred_cnt)]
- ite_list = sorted(ite_list, key=lambda x: x[1], reverse=True)
- n_positive = 0
- idx = 0
- while n_positive < size_cascade:
- if ite_list[idx][2] == 1:
- n_positive += 1
- pre = n_positive / (idx + 1.)
- avgp = avgp + pre
- idx += 1
- # if idx >= 1:
- # break
- avgp = avgp / size_cascade
- return avgp
-
-
- def main():
- '''Main Function'''
-
- parser = argparse.ArgumentParser(description='translate.py')
-
- parser.add_argument('-model', default="/media/external_3TB/3TB/ramezani/pmoini/tmp_debug_NDM/NeuralDiffusionModel-master/pa-our-twitterLastfm_test49.chkpt",
- help='Path to model .pt file')
- '''
- parser.add_argument('-src', required=True,
- help='Source sequence to decode (one line per sequence)')
- parser.add_argument('-vocab', required=True,
- help='Source sequence to decode (one line per sequence)')
- '''
- parser.add_argument('-output', default='pred.txt',
- help="""Path to output the predictions (each line will
- be the decoded sequence""")
- parser.add_argument('-beam_size', type=int, default=100,
- help='Beam size')
- parser.add_argument('-batch_size', type=int, default=1,
- help='Batch size')
- parser.add_argument('-n_best', type=int, default=1,
- help="""If verbose is set, will output the n_best
- decoded sentences""")
- parser.add_argument('-no_cuda', action='store_true')
- # parser.add_argument('-seq_len', type=int ,default=20)
- parser.add_argument('-pred_len', type=int, default=15)
-
- opt = parser.parse_args()
- opt.cuda = not opt.no_cuda
-
- # Prepare DataLoader
-
- test_data = DataLoader(use_valid=False, batch_size=opt.batch_size, cuda=opt.cuda)
-
- translator = Translator(opt)
- translator.model.eval()
-
- numuser = test_data.user_size
-
- num_right = 0
- num_total = 0
-
- avgF1 = 0
- avgPre = 0
- avgRec = 0
-
- avgF1_long = 0
- avgPre_long = 0
- avgRec_long = 0
-
- avgF1_short = 0
- avgPre_short = 0
- avgRec_short = 0
- numseq = 0 # number of test seqs
-
- # for micro pre rec f1
- right = 0.
- pred = 0.
- total = 0.
- right_long = 0.
- pred_long = 0.
- total_long = 0.
- right_short = 0.
- pred_short = 0.
- total_short = 0.
-
- with open(opt.output, 'w') as f:
- with open(idx2vec_addr, 'rb') as handle:
- idx2vec = pickle.load(handle)
- for (src, src_mask, src_lengths, trg, trg_y, trg_mask, trg_lengths) in tqdm(test_data, mininterval=2,
- desc=' - (Test)', leave=False):
- # src = torch.FloatTensor([[ idx2vec[int(idx.data.cpu().numpy())] for idx in src[i]] for i in range(src.size(0))]).cuda(0)
- all_samples = translator.translate_batch(src, src_mask, src_lengths, trg, trg_mask, trg_lengths).data
- for bid in range(src.size(0)):
- numseq += 1.0
-
- ground_truth = np.zeros([numuser])
- num_ground_truth = 0
- for user in trg_y.data[bid][1:-1]:
- if user == Constants.EOS or user == Constants.PAD:
- break
- ground_truth[user] = 1.0
- num_ground_truth += 1
- pred_cnt = np.zeros([numuser])
- for beid in range(opt.beam_size):
- # print(all_samples[0,0,0:1])
- for j in range(num_ground_truth):
- # print(bid,beid,j)
- pred_uid = all_samples[bid, beid, j]
- if pred_uid == Constants.EOS:
- break
- else:
- pred_cnt[pred_uid] += 1.0 / opt.beam_size
-
- F1, pre, rec = getF1(ground_truth, pred_cnt)
- avgF1 += F1
- avgPre += pre
- avgRec += rec
- right += np.dot(ground_truth, pred_cnt)
- pred += np.sum(pred_cnt)
- total += np.sum(ground_truth)
-
- # for short user
- ground_truth = np.zeros([numuser])
- num_ground_truth = 0
- for user in trg_y.data[bid][1:-1]:
- if user == Constants.EOS or user == Constants.PAD:
- break
- ground_truth[user] = 1.0
- num_ground_truth += 1
- if num_ground_truth >= 5:
- break
-
- pred_cnt = np.zeros([numuser])
- for beid in range(opt.beam_size):
- # total += len(ground_truth)
- for j in range(num_ground_truth):
- pred_uid = all_samples[bid, beid, j]
- if pred_uid == Constants.EOS:
- break
- # continue
- else:
- pred_cnt[pred_uid] += 1.0 / opt.beam_size
-
- F1, pre, rec = getF1(ground_truth, pred_cnt)
- avgF1_short += F1
- avgPre_short += pre
- avgRec_short += rec
- right_short += np.dot(ground_truth, pred_cnt)
- pred_short += np.sum(pred_cnt)
- total_short += np.sum(ground_truth)
- print(avgF1, avgPre, avgRec)
-
- print('[Info] Finished.')
- print('Macro')
- print(avgF1 / numseq)
- print(avgPre / numseq)
- print(avgRec / numseq)
- print('Results for the first no more than 5 predictions')
- print(avgF1_short / numseq)
- print(avgPre_short / numseq)
- print(avgRec_short / numseq)
-
- print('Micro')
- pmi = right / pred
- rmi = right / total
- print(2 * pmi * rmi / (pmi + rmi))
- print(pmi)
- print(rmi)
-
- print('Results for the first no more than 5 predictions')
- pmi_long = right_short / pred_short
- rmi_long = right_short / total_short
- print(2 * pmi_long * rmi_long / (pmi_long + rmi_long))
- print(pmi_long)
- print(rmi_long)
-
-
- if __name__ == "__main__":
- main()
|