diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index 2a67e62..2015cf2 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -3,7 +3,17 @@
   <component name="ChangeListManager">
     <list default="true" id="0f79fb12-bf2b-474d-8da4-e622edb018dd" name="Default Changelist" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/Config.py" beforeDir="false" afterPath="$PROJECT_DIR$/Config.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/base/Base.cpp" beforeDir="false" afterPath="$PROJECT_DIR$/base/Base.cpp" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/base/Setting.h" beforeDir="false" afterPath="$PROJECT_DIR$/base/Setting.h" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/base/Test.h" beforeDir="false" afterPath="$PROJECT_DIR$/base/Test.h" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/distribute_training.py" beforeDir="false" afterPath="$PROJECT_DIR$/distribute_training.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/main_spark.py" beforeDir="false" afterPath="$PROJECT_DIR$/main_spark.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/release/Base.so" beforeDir="false" afterPath="$PROJECT_DIR$/release/Base.so" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/run_dbpedia.sh" beforeDir="false" afterPath="$PROJECT_DIR$/run_dbpedia.sh" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/split/generate.py" beforeDir="false" afterPath="$PROJECT_DIR$/split/generate.py" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/test.py" beforeDir="false" afterPath="$PROJECT_DIR$/test.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/test_1.py" beforeDir="false" afterPath="$PROJECT_DIR$/test_1.py" afterDir="false" />
     </list>
     <option name="EXCLUDED_CONVERTED_TO_IGNORED" value="true" />
     <option name="SHOW_DIALOG" value="false" />
@@ -21,7 +31,7 @@
   <component name="ProjectId" id="1On7TJBKj3ZnByQP1GDdyVRo18W" />
   <component name="PropertiesComponent">
     <property name="SHARE_PROJECT_CONFIGURATION_FILES" value="true" />
-    <property name="last_opened_file_path" value="$USER_HOME$/Desktop/DBpedia/generate.py" />
+    <property name="last_opened_file_path" value="$PROJECT_DIR$" />
     <property name="project.structure.last.edited" value="Problems" />
     <property name="project.structure.proportion" value="0.15" />
     <property name="project.structure.side.proportion" value="0.2" />
@@ -30,6 +40,7 @@
   <component name="RecentsManager">
     <key name="CopyFile.RECENT_KEYS">
       <recent name="$PROJECT_DIR$" />
+      <recent name="$PROJECT_DIR$/" />
     </key>
   </component>
   <component name="RunDashboard">
@@ -45,6 +56,48 @@
     </option>
   </component>
   <component name="RunManager" selected="Python.test_1">
+    <configuration name="generate (1)" type="PythonConfigurationType" factoryName="Python" temporary="true">
+      <module name="OpenKEonSpark" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/generate.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
+    <configuration name="generate" type="PythonConfigurationType" factoryName="Python" temporary="true">
+      <module name="OpenKEonSpark" />
+      <option name="INTERPRETER_OPTIONS" value="" />
+      <option name="PARENT_ENVS" value="true" />
+      <envs>
+        <env name="PYTHONUNBUFFERED" value="1" />
+      </envs>
+      <option name="SDK_HOME" value="" />
+      <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/split" />
+      <option name="IS_MODULE_SDK" value="true" />
+      <option name="ADD_CONTENT_ROOTS" value="true" />
+      <option name="ADD_SOURCE_ROOTS" value="true" />
+      <option name="SCRIPT_NAME" value="$PROJECT_DIR$/split/generate.py" />
+      <option name="PARAMETERS" value="" />
+      <option name="SHOW_COMMAND_LINE" value="false" />
+      <option name="EMULATE_TERMINAL" value="false" />
+      <option name="MODULE_MODE" value="false" />
+      <option name="REDIRECT_INPUT" value="false" />
+      <option name="INPUT_FILE" value="" />
+      <method v="2" />
+    </configuration>
     <configuration name="test" type="PythonConfigurationType" factoryName="Python" temporary="true">
       <module name="OpenKEonSpark" />
       <option name="INTERPRETER_OPTIONS" value="" />
@@ -90,6 +143,8 @@
     <recent_temporary>
       <list>
         <item itemvalue="Python.test_1" />
+        <item itemvalue="Python.generate (1)" />
+        <item itemvalue="Python.generate" />
         <item itemvalue="Python.test" />
       </list>
     </recent_temporary>
diff --git a/Config.py b/Config.py
index 5dd4e06..e74f180 100644
--- a/Config.py
+++ b/Config.py
@@ -1,5 +1,6 @@
 #coding:utf-8
 import numpy as np
+import matplotlib.pyplot as plt
 import tensorflow as tf
 import os
 import ctypes
@@ -44,6 +45,12 @@ def __init__(self, cpp_lib_path=None, init_new_entities=False):
             self.lib.getBestThreshold.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
             self.lib.test_triple_classification.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
 
+            #ROC
+            self.lib.get_n_interval.argtypes = [ctypes.c_int64, ctypes.c_void_p, ctypes.c_void_p]
+            self.lib.get_n_interval.restype = ctypes.c_int64
+            self.lib.get_TPFP.argtypes = [ctypes.c_int64, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p]
+            self.lib.get_TPFP.restype = ctypes.POINTER( ctypes.c_int64 * 2 )
+
             #set other parameters
             self.in_path = None
             self.test_log_path = None
@@ -54,7 +61,7 @@ def __init__(self, cpp_lib_path=None, init_new_entities=False):
             self.rel_size = self.hidden_size
             self.train_times = 0
             self.margin = 1.0
-            self.nbatches = 100
+            self.nbatches = 0
             self.negative_ent = 1
             self.negative_rel = 0
             self.workThreads = 8
@@ -852,6 +859,60 @@ def test(self):
                 print("\nElapsed test time (seconds): {}".format(test_time_elapsed))
 
 
+    def plot_roc(self, rel_index, fig_name=None):
+        if self.importName != None:
+            self.restore_tensorflow()
+        self.init_triple_classification()
+
+        self.lib.getValidBatch(self.valid_pos_h_addr, self.valid_pos_t_addr, self.valid_pos_r_addr, self.valid_neg_h_addr, self.valid_neg_t_addr, self.valid_neg_r_addr)
+        res_pos_valid = self.test_step(self.valid_pos_h, self.valid_pos_t, self.valid_pos_r)
+        res_neg_valid = self.test_step(self.valid_neg_h, self.valid_neg_t, self.valid_neg_r)
+
+        self.lib.getTestBatch(self.test_pos_h_addr, self.test_pos_t_addr, self.test_pos_r_addr, self.test_neg_h_addr, self.test_neg_t_addr, self.test_neg_r_addr)
+        res_pos_test = self.test_step(self.test_pos_h, self.test_pos_t, self.test_pos_r)
+        res_neg_test = self.test_step(self.test_neg_h, self.test_neg_t, self.test_neg_r)
+
+        n_intervals = self.lib.get_n_interval(rel_index, res_pos_valid.__array_interface__['data'][0], res_neg_valid.__array_interface__['data'][0])
+        self.lib.get_TPFP.restype = ctypes.POINTER( ctypes.c_int64 * ((n_intervals+1)*2) )
+        res = [j for j in self.lib.get_TPFP(rel_index, res_pos_valid.__array_interface__['data'][0], res_neg_valid.__array_interface__['data'][0], res_pos_test.__array_interface__['data'][0], res_neg_test.__array_interface__['data'][0]).contents]
+
+        TPR = []
+        FPR = []
+
+        if res[0] != 0 or res[0+n_intervals+1] != 0:
+            TPR.append(0)
+            FPR.append(0)
+
+
+        for i in range(0, n_intervals+1):
+            TPR.append(res[i])
+            FPR.append(res[i+n_intervals+1])
+
+        if TPR[len(TPR)-1] != len(res_pos_test.flatten()) or FPR[len(FPR)-1] != len(res_neg_test.flatten()):
+            TPR.append(len(res_pos_test.flatten()))
+            FPR.append(len(res_neg_test.flatten()))
+
+
+        for i in range(len(TPR)): TPR[i] /= TPR[-1]
+        for i in range(len(FPR)): FPR[i] /= FPR[-1]
+
+        auc = np.trapz(TPR, FPR)
+
+        plt.figure()
+        lw=2
+        plt.plot(FPR, TPR, color='darkorange', lw=lw, label='ROC curve (area = %0.3f)' % auc)
+        plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.05])
+        plt.xlabel('False Positive Rate (FPR)')
+        plt.ylabel('True Positive Rate (TPR)')
+        plt.title('ROC Curve')
+        plt.legend(loc="lower right")
+        if fig_name == None or fig_name == '':
+            plt.show()
+        else:
+            plt.savefig(fig_name)
+
 
     def predict_head_entity(self, t, r, k):
         r'''This mothod predicts the top k head entities given tail entity and relation.
diff --git a/__pycache__/Config.cpython-36.pyc b/__pycache__/Config.cpython-36.pyc
index c38d694..94c72a7 100644
Binary files a/__pycache__/Config.cpython-36.pyc and b/__pycache__/Config.cpython-36.pyc differ
diff --git a/__pycache__/TransD.cpython-36.pyc b/__pycache__/TransD.cpython-36.pyc
index 9bf4020..c765123 100644
Binary files a/__pycache__/TransD.cpython-36.pyc and b/__pycache__/TransD.cpython-36.pyc differ
diff --git a/__pycache__/TransE.cpython-36.pyc b/__pycache__/TransE.cpython-36.pyc
index 46528fe..f7e444a 100644
Binary files a/__pycache__/TransE.cpython-36.pyc and b/__pycache__/TransE.cpython-36.pyc differ
diff --git a/__pycache__/distribute_training.cpython-36.pyc b/__pycache__/distribute_training.cpython-36.pyc
index 59a2ad0..05ef2fe 100644
Binary files a/__pycache__/distribute_training.cpython-36.pyc and b/__pycache__/distribute_training.cpython-36.pyc differ
diff --git a/base/Base.cpp b/base/Base.cpp
index e0ddb16..3417a7f 100644
--- a/base/Base.cpp
+++ b/base/Base.cpp
@@ -55,6 +55,8 @@ void importTrainFiles();
 extern "C"
 void importOntologyFiles();
 
+
+
 struct Parameter {
 	INT id;
 	INT *batch_h;
diff --git a/base/Setting.h b/base/Setting.h
index 5c8504f..2df9cf5 100644
--- a/base/Setting.h
+++ b/base/Setting.h
@@ -111,6 +111,10 @@ extern "C"
 void setBern(INT con) {
 	bernFlag = con;
 }
+/*
+============================================================
+*/
+REAL interval = 0.01;
 
 
 #endif
diff --git a/base/Test.h b/base/Test.h
index 501715f..7ea04dd 100644
--- a/base/Test.h
+++ b/base/Test.h
@@ -302,9 +302,9 @@ void getValidBatch(INT *ph, INT *pt, INT *pr, INT *nh, INT *nt, INT *nr) {
 REAL threshEntire;
 extern "C"
 void getBestThreshold(REAL *relThresh, REAL *score_pos, REAL *score_neg) {
-    REAL interval = 0.01;
     REAL min_score, max_score, bestThresh, tmpThresh, bestAcc, tmpAcc;
     INT n_interval, correct, total;
+
     for (INT r = 0; r < relationTotal; r++) {
         if (validLef[r] == -1) continue;
         total = (validRig[r] - validLef[r] + 1) * 2;
@@ -319,6 +319,7 @@ void getBestThreshold(REAL *relThresh, REAL *score_pos, REAL *score_neg) {
             if(score_neg[i] > max_score) max_score = score_neg[i];
         }
         n_interval = INT((max_score - min_score)/interval);
+
         for (INT i = 0; i <= n_interval; i++) {
             tmpThresh = min_score + i * interval;
             correct = 0;
@@ -345,8 +346,8 @@ extern "C"
 //EDIT
 void test_triple_classification(REAL *relThresh, REAL *score_pos, REAL *score_neg, REAL *acc_addr) {
     testAcc = (REAL *)calloc(relationTotal, sizeof(REAL));
-    INT TP = 0, TN = 0, FP = 0, FN = 0;
     REAL accuracy, precision, recall, fmeasure;
+    INT TP = 0, TN = 0, FP = 0, FN = 0;
 
     for (INT r = 0; r < relationTotal; r++) {
         if (validLef[r] == -1 || testLef[r] ==-1) continue;
@@ -386,5 +387,62 @@ void test_triple_classification(REAL *relThresh, REAL *score_pos, REAL *score_ne
 }
 
 
+extern "C"
+INT get_n_interval(INT r, REAL *score_pos, REAL *score_neg){
+    REAL min_score, max_score;
+    INT n_interval, total;
+    if (validLef[r] == -1) return 0;
+    total = (validRig[r] - validLef[r] + 1) * 2;
+    min_score = score_pos[validLef[r]];
+    if (score_neg[validLef[r]] < min_score) min_score = score_neg[validLef[r]];
+    max_score = score_pos[validLef[r]];
+    if (score_neg[validLef[r]] > max_score) max_score = score_neg[validLef[r]];
+    for (INT i = validLef[r]+1; i <= validRig[r]; i++) {
+        if(score_pos[i] < min_score) min_score = score_pos[i];
+        if(score_pos[i] > max_score) max_score = score_pos[i];
+        if(score_neg[i] < min_score) min_score = score_neg[i];
+        if(score_neg[i] > max_score) max_score = score_neg[i];
+    }
+    return INT((max_score - min_score)/interval);
+}
+
+
+extern "C"
+INT* get_TPFP(INT r, REAL *score_pos, REAL *score_neg, REAL *score_pos_test, REAL *score_neg_test) {
+    REAL min_score, max_score, tmpThresh;
+    INT n_interval, total;
+
+
+    if (validLef[r] == -1) return 0;
+    total = (validRig[r] - validLef[r] + 1) * 2;
+    min_score = score_pos[validLef[r]];
+    if (score_neg[validLef[r]] < min_score) min_score = score_neg[validLef[r]];
+    max_score = score_pos[validLef[r]];
+    if (score_neg[validLef[r]] > max_score) max_score = score_neg[validLef[r]];
+    for (INT i = validLef[r]+1; i <= validRig[r]; i++) {
+        if(score_pos[i] < min_score) min_score = score_pos[i];
+        if(score_pos[i] > max_score) max_score = score_pos[i];
+        if(score_neg[i] < min_score) min_score = score_neg[i];
+        if(score_neg[i] > max_score) max_score = score_neg[i];
+    }
+    n_interval = INT((max_score - min_score)/interval);
+
+
+    INT* TPFPs = new INT[(n_interval+1)*2];
+    for (INT i = 0; i <= n_interval; i++) {
+        INT TP = 0, FP = 0;
+        tmpThresh = min_score + i * interval;
+        for (INT i = testLef[r]; i <= testRig[r]; i++) {
+            if (score_pos_test[i] <= tmpThresh) TP++;
+            if (score_neg_test[i] <= tmpThresh) FP++;
+        }
+        TPFPs[i] = TP;
+        TPFPs[i + n_interval+1] = FP;
+    }
+
+    return TPFPs;
+}
+
+
 
 #endif
diff --git a/distribute_training.py b/distribute_training.py
index 0ff3d7d..ffd4d40 100644
--- a/distribute_training.py
+++ b/distribute_training.py
@@ -172,7 +172,7 @@ def main_fun(argv, ctx):
 
         if sys.argv.debug: print("Creating Hooks, Scaffold, FileWriter, ConfigProto...")
         iterations = con.train_times * con.nbatches + last_global_step
-        hooks=[tf.train.StopAtStepHook(last_step=iterations)]
+        # hooks=[tf.train.StopAtStepHook(last_step=iterations)]
         scaffold = tf.train.Scaffold(init_op=init_op, saver=saver, summary_op=summary_op)
         tf.summary.FileWriter("tensorboard_%d" % ctx.worker_num, graph=tf.get_default_graph())
         config_monitored = tf.ConfigProto()
@@ -196,7 +196,7 @@ def main_fun(argv, ctx):
                                                checkpoint_dir=argv.output_path,
                                                save_summaries_steps=1,
                                                save_checkpoint_steps=1*con.nbatches,
-                                               hooks=hooks,
+                                               # hooks=hooks,
                                                summary_dir="tensorboard_%d" % ctx.worker_num
                                                )
         if sys.argv.debug:
@@ -215,99 +215,113 @@ def main_fun(argv, ctx):
         to_reach_step = sys.argv.early_stop_start_step * con.nbatches + last_global_step
         con.lib.getValidBatch(con.valid_pos_h_addr, con.valid_pos_t_addr, con.valid_pos_r_addr, con.valid_neg_h_addr, con.valid_neg_t_addr, con.valid_neg_r_addr)
         local = 0
+        g = last_global_step
 
-        while not sess.should_stop():
-            #gives time to other workers to connect
-            if task_index == 0 and local == 0:
-                time.sleep(60)
-            local += 1
 
-            con.sampling()
+        while g < iterations:
+            try:
+                #gives time to other workers to connect
+                if task_index == 0 and local == 0:
+                    time.sleep(30)
 
-            feed_dict = {
-                trainModel.batch_h: con.batch_h,
-                trainModel.batch_t: con.batch_t,
-                trainModel.batch_r: con.batch_r,
-                trainModel.batch_y: con.batch_y
-            }
+                local += 1
+                con.sampling()
+                feed_dict = {
+                    trainModel.batch_h: con.batch_h,
+                    trainModel.batch_t: con.batch_t,
+                    trainModel.batch_r: con.batch_r,
+                    trainModel.batch_y: con.batch_y
+                }
 
-            _, loss, g = sess.run([train_op, trainModel.loss, global_step], feed_dict)
-            print('Global step: {} Epoch: {} Batch: {} loss: {}'.format(g, int( (g-last_global_step) /con.nbatches), int( (g - last_global_step) % con.nbatches), loss))
+                _, loss, g = sess.run([train_op, trainModel.loss, global_step], feed_dict)
+                print('Global step: {} Epoch: {} Batch: {} loss: {}'.format(g, int( (g-last_global_step) /con.nbatches), int( (g - last_global_step) % con.nbatches), loss))
 
 
-            ################## EARLY STOP ##################
-            if (task_index != 0) and (g >= to_reach_step):
-                to_reach_step += stopping_step
-                if os.path.exists(sys.argv.output_path+"/stop.txt"):
-                    print('\nEarly stop happened in chief worker\n')
-                    break
+                ################## EARLY STOP ##################
+                if (task_index != 0) and (g >= to_reach_step):
+                    to_reach_step += stopping_step
+                    if os.path.exists(sys.argv.output_path+"/stop.txt"):
+                        print('\nEarly stop happened in chief worker\n')
+                        break
 
 
-            if (task_index == 0) and (not sess.should_stop()) and (g >= to_reach_step):
-                while (g >= to_reach_step):
-                    to_reach_step += stopping_step
+                if (task_index == 0) and (g < iterations) and (g >= to_reach_step):
+                    while (g >= to_reach_step):
+                        to_reach_step += stopping_step
+
+
+                    ################## ACCURACY ##################
+                    feed_dict[trainModel.predict_h] = con.valid_pos_h
+                    feed_dict[trainModel.predict_t] = con.valid_pos_t
+                    feed_dict[trainModel.predict_r] = con.valid_pos_r
+                    res_pos = sess.run(trainModel.predict, feed_dict)
+
+                    feed_dict[trainModel.predict_h] = con.valid_neg_h
+                    feed_dict[trainModel.predict_t] = con.valid_neg_t
+                    feed_dict[trainModel.predict_r] = con.valid_neg_r
+                    res_neg = sess.run(trainModel.predict, feed_dict)
+
+
+                    con.lib.getBestThreshold(con.relThresh_addr, res_pos.__array_interface__['data'][0], res_neg.__array_interface__['data'][0])
+                    con.lib.test_triple_classification(con.relThresh_addr, res_pos.__array_interface__['data'][0], res_neg.__array_interface__['data'][0], con.acc_addr)
+
+
+                    acc = con.acc[0]
+
+                    if sys.argv.debug:
+                        print("\n[ Early Stop Check (Accuracy) ]")
+                        print("Best Accuracy = %.10f" %(best_acc))
+                        print("Accuracy after run  =  %.10f" %(acc))
+
+                    if acc > best_acc:
+                        best_acc = acc
+                        wait_steps_acc = 0
+                        if sys.argv.debug: print("New best Accuracy founded. Wait steps reset.")
+                        best_model_global_step_acc = g
+
+                    elif wait_steps_acc < patience:
+                        wait_steps_acc += 1
+                        if sys.argv.debug: print("Wait steps Accuracy incremented: {}\n".format(wait_steps_acc))
+
+
+                    if wait_steps_acc >= patience:
+                        print('Accuracy early stop. Accuracy has not been improved enough in {} times'.format(patience))
+                        with open(sys.argv.output_path+"/stop.txt", "w") as stop_file:
+                            stop_file.write(str(best_model_global_step_acc)+"\n")
+                        break
+
+                    ################## LOSS ##################
+                    if sys.argv.debug:
+                        print("\n[ Early Stop Checking (Loss) ]")
+                        print("Best loss = %.10f" %(best_loss))
+                        print("Loss after run  =  %.10f" %(loss))
+
+                    if loss < best_loss:
+                        best_loss = loss
+                        wait_steps_loss = 0
+                        if sys.argv.debug: print("New best loss founded. Wait steps reset.")
+                        best_model_global_step_loss = g
+
+                    elif wait_steps_loss < patience:
+                        wait_steps_loss += 1
+                        if sys.argv.debug: print("Wait steps loss incremented: {}\n".format(wait_steps_loss))
+
+
+                    if wait_steps_loss >= patience:
+                        print('Loss early stop. Losses has not been improved enough in {} times'.format(patience))
+                        with open(sys.argv.output_path+"/stop.txt", "w") as stop_file:
+                            stop_file.write(str(best_model_global_step_loss)+"\n")
+                        break
+
+            except:
+                print("Exception occured; stopping training")
+                break
+
+        #gives time to chief to stop
+        if task_index != 0:
+            time.sleep(15)
+
 
-                ################## ACCURACY ##################
-                feed_dict[trainModel.predict_h] = con.valid_pos_h
-                feed_dict[trainModel.predict_t] = con.valid_pos_t
-                feed_dict[trainModel.predict_r] = con.valid_pos_r
-                res_pos = sess.run(trainModel.predict, feed_dict)
-
-                feed_dict[trainModel.predict_h] = con.valid_neg_h
-                feed_dict[trainModel.predict_t] = con.valid_neg_t
-                feed_dict[trainModel.predict_r] = con.valid_neg_r
-                res_neg = sess.run(trainModel.predict, feed_dict)
-
-                con.lib.getBestThreshold(con.relThresh_addr, res_pos.__array_interface__['data'][0], res_neg.__array_interface__['data'][0])
-                con.lib.test_triple_classification(con.relThresh_addr, res_pos.__array_interface__['data'][0], res_neg.__array_interface__['data'][0], con.acc_addr)
-                acc = con.acc[0]
-
-                if sys.argv.debug:
-                    print("\n[ Early Stop Check (Accuracy) ]")
-                    print("Best Accuracy = %.10f" %(best_acc))
-                    print("Accuracy after run  =  %.10f" %(acc))
-
-                if acc > best_acc:
-                    best_acc = acc
-                    wait_steps_acc = 0
-                    if sys.argv.debug: print("New best Accuracy founded. Wait steps reset.")
-                    best_model_global_step_acc = g
-
-                elif wait_steps_acc < patience:
-                    wait_steps_acc += 1
-                    if sys.argv.debug: print("Wait steps Accuracy incremented: {}\n".format(wait_steps_acc))
-
-
-                if wait_steps_acc >= patience:
-                    print('Accuracy early stop. Accuracy has not been improved enough in {} times'.format(patience))
-                    with open(sys.argv.output_path+"/stop.txt", "w") as stop_file:
-                        stop_file.write(str(best_model_global_step_acc)+"\n")
-                    break
-
-                ################## LOSS ##################
-                if sys.argv.debug:
-                    print("\n[ Early Stop Checking (Loss) ]")
-                    print("Best loss = %.10f" %(best_loss))
-                    print("Loss after run  =  %.10f" %(loss))
-
-                if loss < best_loss:
-                    best_loss = loss
-                    wait_steps_loss = 0
-                    if sys.argv.debug: print("New best loss founded. Wait steps reset.")
-                    best_model_global_step_loss = g
-
-                elif wait_steps_loss < patience:
-                    wait_steps_loss += 1
-                    if sys.argv.debug: print("Wait steps loss incremented: {}\n".format(wait_steps_loss))
-
-
-                if wait_steps_loss >= patience:
-                    print('Loss early stop. Losses has not been improved enough in {} times'.format(patience))
-                    with open(sys.argv.output_path+"/stop.txt", "w") as stop_file:
-                        stop_file.write(str(best_model_global_step_loss)+"\n")
-                    break
-
-        return
 
 
 
diff --git a/main_spark.py b/main_spark.py
index 611e6fc..fa73ce1 100644
--- a/main_spark.py
+++ b/main_spark.py
@@ -342,7 +342,7 @@ def n_n():
 
 
     if args.debug: print("Shutdowning cluster...")
-    cluster.shutdown()
+    cluster.shutdown(timeout=-1)
     training_time = time.time() - training_time
 
 
diff --git a/release/Base.so b/release/Base.so
index 2b8e904..3017fe8 100755
Binary files a/release/Base.so and b/release/Base.so differ
diff --git a/run_dbpedia.sh b/run_dbpedia.sh
index ca7faec..4129e27 100644
--- a/run_dbpedia.sh
+++ b/run_dbpedia.sh
@@ -69,7 +69,7 @@ do
     --cluster_size $SPARK_WORKER_INSTANCES --num_ps 1 --num_gpus 1 --cpp_lib_path $WORK_DIR_PREFIX/release/Base.so \
 	--input_path /content/drive/My\ Drive/DBpedia/$n/$i/ \
     --output_path $WORK_DIR_PREFIX/res_spark \
-    --alpha $4 --optimizer SGD --train_times 50 --ent_neg_rate 1 --embedding_dimension $2 --margin 1.0 --model $3
+    --alpha $4 --optimizer SGD --train_times 50 --ent_neg_rate 1 --embedding_dimension $2 --margin 1.0 --model $3 --n_batches 100
 
 
 	echo "====================================== Copying model for batch $i ======================================"
diff --git a/split/generate.py b/split/generate.py
index 7fe9288..3400942 100644
--- a/split/generate.py
+++ b/split/generate.py
@@ -16,10 +16,10 @@
 TARGET_RELATION = '<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>'
 
 #dataset containing all the triples with only the target relation
-path_t = "instancetype_en.nt"
+path_t = "/home/luigi/Desktop/DBpedia/instancetype_en.nt"
 
 #dataset containing all the other triples (except the target relation)
-path_r = "newstrictmappingbased_en.nt"
+path_r = "/home/luigi/Desktop/DBpedia/newstrictmappingbased_en.nt"
 
 #path which contains the class hierarchy
 path_h = "class_hierarchy.txt"
@@ -35,7 +35,7 @@
 GENERATE_ONTOLOGY_FILES = True
 
 #number of batches to split the dataset into
-N_BATCHES = 5
+N_BATCHES = 1
 
 #percentage of triples with target relation for each batch test set 
 TEST_SET_PERCENTAGE = 10
@@ -59,7 +59,7 @@
 type_tails_dict = {}
 
 
-print(" ===== Reading test triples ===== ")
+print(" ===== Reading TARGET triples ===== ")
 with open(path_t, "r") as f:
     lines_t = f.readlines()
     
diff --git a/test.py b/test.py
index aaec590..d6be97d 100644
--- a/test.py
+++ b/test.py
@@ -9,11 +9,11 @@
     print(type(arg), arg)
     print("\n")
 
-n = sys.argv[1]
-max = sys.argv[2]
-dim = sys.argv[3]
-model = sys.argv[4]
-lp = sys.argv[5]
+path_to_append = sys.argv[1]
+# max = sys.argv[2]
+dim = sys.argv[2]
+model = sys.argv[3]
+lp = sys.argv[4]
 
 
 
@@ -25,7 +25,7 @@ def get_ckpt(p):
         ckpt = ckpt[len(ckpt) - 1]
     return ckpt
 
-dataset_path = '/content/drive/My Drive/DBpedia/{}/{}/'.format(max, n)
+dataset_path = '/content/drive/My Drive/DBpedia/{}'.format(path_to_append)
 path = dataset_path + 'model/'
 ckpt = get_ckpt(path)
 
@@ -35,18 +35,23 @@ def get_ckpt(p):
 con.set_test_triple_classification(True)
 con.set_dimension(int(dim))
 con.init()
-con.set_n_threads_LP(2)
+
 
 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')
diff --git a/test_1.py b/test_1.py
index 0abee2d..0e59c5b 100644
--- a/test_1.py
+++ b/test_1.py
@@ -4,8 +4,8 @@
 from TransD import TransD
 from TransR import TransR
 import sys
-# import os
-# os.environ["CUDA_VISIBLE_DEVICES"]="-1"
+
+
 
 def get_ckpt(p):
     ckpt = None
@@ -17,7 +17,8 @@ def get_ckpt(p):
 
 
 #/home/luigi/IdeaProjects/OpenKE_new_Spark/benchmarks/DBpedia
-dataset_path = '/home/luigi/files/stuff/DBpedia/5/0/'
+# dataset_path = '/content/drive/My Drive/DBpedia/1/0/'
+dataset_path = '/home/luigi/files/stuff/DBpedia/5/1/'
 # dataset_path = '/home/luigi/files/stuff/superuser/9/1/'
 path = dataset_path + 'model/'
 # path = '/home/luigi/IdeaProjects/OpenKEonSpark/res_spark/'
@@ -25,20 +26,18 @@ def get_ckpt(p):
 ckpt = get_ckpt(path)
 
 
+# con = Config(cpp_lib_path='/content/OpenKEonSpark/release/Base.so')
 con = Config()
 con.set_in_path(dataset_path)
-con.set_test_link_prediction(True)
-con.set_test_triple_classification(True)
+# con.set_test_link_prediction(True)
+# con.set_test_triple_classification(True)
 con.set_dimension(int(64))
 con.init()
-con.set_model_and_session(TransD)
+con.set_model_and_session(TransE)
 con.set_import_files(path+ckpt)
-con.set_test_log_path(path)
-con.set_n_threads_LP(5)
-con.test()
-
-con.predict_tail_entity(349585, 5, 10)
-# for i in range(0,100):
-#     con.predict_tail_entity(i,0,1)
-# print(con.acc)
+# con.set_test_log_path(path)
+# con.set_n_threads_LP(1)
+# con.test()
+con.plot_roc(5)
+