author net predict

Author: NiMaZi

predict.py

@@ -21,8 +21,9 @@
 parser.add_argument('-vocab',default='data_sample/vocab/Vid2Name.json',type=str,help="entity dictionary")
 parser.add_argument('-A',default=False,type=bool,help="train with author info")
 parser.add_argument('-E',default=False,type=bool,help="evaluate with real body")
-parser.add_argument('-abstract',default='data_sample/sample/abs0.csv',type=str,help="abstract to predict")
-parser.add_argument('-body',default='data_sample/sample/body0.csv',type=str,help="body to evaluate")
+parser.add_argument('-abstract',default='data_sample/sample_author/abs0.csv',type=str,help="abstract to predict")
+parser.add_argument('-authors',default='data_sample/sample_author/authors0.json',type=str,help="author information")
+parser.add_argument('-body',default='data_sample/sample_author/body0.csv',type=str,help="body to evaluate")
 parser.add_argument('-path',default='model.h5',type=str,help="path for loading the model")
 parser.add_argument('-threshold',default=0.0,type=float,help="threshold of prediction, 0.0~1.0")
 opt=parser.parse_args()
@@ -51,8 +52,39 @@ def get_prediction(model,abs_path,in_dict):
 	return abs_vec,list(pred)
 
 
-def get_prediction_author(model,abs_path,in_dict,author_dict):
-	pass
+def get_prediction_author(model,abs_path,authors_path,in_dict,author_dict):
+	cc2vid_input=in_dict
+	fa2vid=author_dict
+
+	abs_vec=[0.0 for i in range(0,len(cc2vid_input))]
+	abs_count=0.0
+
+	with open(abs_path,'r',encoding='utf-8') as cf:
+		rd=csv.reader(cf)
+		for item in rd:
+			if item[0]=="Mention":
+				continue
+			try:
+				abs_vec[cc2vid_input[item[1]]]+=1.0
+				abs_count+=1.0
+			except:
+				pass
+	if not abs_count==0.0:
+		abs_vec=list(np.array(abs_vec)/abs_count)
+
+	author_vec=[0.0 for i in range(0,len(fa2vid))]
+	authors=util.load_sups(authors_path)
+	for author in authors:
+		try:
+			author_vec[fa2vid[author]]=1.0
+		except:
+			pass
+
+	sample_input=np.array([abs_vec+author_vec])
+	pred=model.model.predict([sample_input[:,:len(cc2vid_input)],sample_input[:,len(cc2vid_input):]])[0]
+	pred/=np.linalg.norm(pred)
+
+	return abs_vec,list(pred)
 
 def print_vec(prediction,entity_dict,threshold=0.0):
 	for i,v in enumerate(prediction):
@@ -120,11 +152,16 @@ def main():
 
 	if not author:
 		bownn_model=BOWNN()
+		bownn_model.load_model(load_path)
+		abs_vec,prediction=get_prediction(bownn_model,abs_path,in_dict)
 	else:
+		author_path=opt.author
+		author_dict=util.load_sups(author_path)
+		authors_path=opt.authors
 		bownn_model=BOWNN_author()
+		bownn_model.load_model(load_path)
+		abs_vec,prediction=get_prediction_author(bownn_model,abs_path,authors_path,in_dict,author_dict)
 
-	bownn_model.load_model(load_path)
-	abs_vec,prediction=get_prediction(bownn_model,abs_path,in_dict)
 	print("Entity mentions in this abstract:")
 	print_vec(abs_vec,entity_dict)
 	print("\n")