Evaluation mode

test.py

@@ -4,54 +4,42 @@
 from TransR import TransR
 from TransD import TransD
 import sys
+import os
+os.environ['CUDA_VISIBLE_DEVICES']='-1'
+for arg in sys.argv: print(arg, type(arg))
+target_rel_index = None
 
-for arg in sys.argv:
-    print(type(arg), arg)
-    print("\n")
-
-path_to_append = sys.argv[1]
-# max = sys.argv[2]
-dim = sys.argv[2]
-model = sys.argv[3]
-lp = sys.argv[4]
+dataset_path = sys.argv[1]
+model_path = sys.argv[2]
+cpp_path = sys.argv[3]
+dim = sys.argv[4]
+model = sys.argv[5]
+if (len(sys.argv) >= 7): target_rel_index = sys.argv[6]
 
 
 
 def get_ckpt(p):
     ckpt = None
-    with open(p + "checkpoint", 'r') as f:
+    with open(os.path.join(p,"checkpoint"), 'r') as f:
         first_line = f.readline()
         ckpt = first_line.split(':')[1].strip().replace('"', '').split('/')
         ckpt = ckpt[len(ckpt) - 1]
     return ckpt
 
-dataset_path = '/content/drive/My Drive/DBpedia/{}'.format(path_to_append)
-path = dataset_path + 'model/'
-ckpt = get_ckpt(path)
 
-con = Config(cpp_lib_path='/content/OpenKEonSpark/release/Base.so')
+ckpt = get_ckpt(model_path)
+con = Config(cpp_lib_path=cpp_path)
 con.set_in_path(dataset_path)
-con.set_test_link_prediction(bool(int(lp)))
 con.set_test_triple_classification(True)
 con.set_dimension(int(dim))
 con.init()
 
+if model.lower() == "transe": con.set_model_and_session(TransE)
+elif model.lower() == "transh": con.set_model_and_session(TransH)
+elif model.lower() == "transr": con.set_model_and_session(TransR)
+else: con.set_model_and_session(TransD)
+
+con.set_import_files(os.path.join(model_path, ckpt))
 
-if model.lower() == "transe":
-    con.set_model_and_session(TransE)
-    con.set_n_threads_LP(5)
-elif model.lower() == "transh":
-    con.set_model_and_session(TransH)
-    con.set_n_threads_LP(5)
-elif model.lower() == "transr":
-    con.set_model_and_session(TransR)
-    con.set_n_threads_LP(2)
-else:
-    con.set_model_and_session(TransD)
-    con.set_n_threads_LP(2)
-
-con.set_import_files(path+ckpt)
-con.set_test_log_path(path)
 con.test()
-print(con.acc)
-con.plot_roc(rel_index=5, fig_name='plot.png')
+if target_rel_index != None: con.plot_roc(rel_index=int(target_rel_index), fig_name='plot.png')