From 29bf2dbcf7cf664818af1ceb54e417e1eebfed78 Mon Sep 17 00:00:00 2001
From: luigiba <baldarigigi@gmail.com>
Date: Fri, 6 Sep 2019 15:51:58 +0200
Subject: [PATCH] Refactoring

---
 .idea/workspace.xml                           |  10 ++
 Config.py                                     |   2 +-
 __pycache__/Config.cpython-36.pyc             | Bin 21054 -> 21017 bytes
 .../distribute_training.cpython-36.pyc        | Bin 15359 -> 15417 bytes
 base/Base.cpp                                 |  11 +-
 base/Reader.h                                 |  12 +-
 base/Setting.h                                |   3 +-
 commands.txt => colab/commands.txt            |  10 +-
 gpu_info.py => colab/gpu_info.py              |   0
 colab/run_dbpedia.sh                          | 125 ++++++++++++++++++
 distribute_training.py                        |   9 +-
 main_spark.py                                 |  29 ++--
 run_dbpedia.sh                                |  87 ------------
 split/generate.py                             |  72 ++++++----
 test.py                                       |  22 ++-
 15 files changed, 236 insertions(+), 156 deletions(-)
 rename commands.txt => colab/commands.txt (72%)
 rename gpu_info.py => colab/gpu_info.py (100%)
 create mode 100644 colab/run_dbpedia.sh
 delete mode 100644 run_dbpedia.sh

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index cedbd43..12ca72b 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -4,8 +4,15 @@
     <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/Reader.h" beforeDir="false" afterPath="$PROJECT_DIR$/base/Reader.h" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/base/Setting.h" beforeDir="false" afterPath="$PROJECT_DIR$/base/Setting.h" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/commands.txt" beforeDir="false" afterPath="$PROJECT_DIR$/colab/commands.txt" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/distribute_training.py" beforeDir="false" afterPath="$PROJECT_DIR$/distribute_training.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/gpu_info.py" beforeDir="false" afterPath="$PROJECT_DIR$/colab/gpu_info.py" afterDir="false" />
       <change beforePath="$PROJECT_DIR$/main_spark.py" beforeDir="false" afterPath="$PROJECT_DIR$/main_spark.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/run_dbpedia.sh" beforeDir="false" afterPath="$PROJECT_DIR$/colab/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" />
     </list>
     <option name="EXCLUDED_CONVERTED_TO_IGNORED" value="true" />
@@ -31,6 +38,9 @@
     <property name="settings.editor.selected.configurable" value="preferences.pluginManager" />
   </component>
   <component name="RecentsManager">
+    <key name="MoveFile.RECENT_KEYS">
+      <recent name="$PROJECT_DIR$/colab" />
+    </key>
     <key name="CopyFile.RECENT_KEYS">
       <recent name="$PROJECT_DIR$" />
       <recent name="$PROJECT_DIR$/" />
diff --git a/Config.py b/Config.py
index d5cec20..25e3cab 100644
--- a/Config.py
+++ b/Config.py
@@ -25,7 +25,7 @@ def __init__(self, cpp_lib_path=None, init_new_entities=False):
         if init_new_entities == False:
             #C library
             if cpp_lib_path == None:
-                cpp_lib_path = '/home/luigi/IdeaProjects/OpenKEonSpark/release/Base.so'
+                cpp_lib_path = './release/Base.so'
             base_file = os.path.abspath(cpp_lib_path)
             self.lib = ctypes.cdll.LoadLibrary(base_file)
             self.lib.sampling.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int64, ctypes.c_int64, ctypes.c_int64]
diff --git a/__pycache__/Config.cpython-36.pyc b/__pycache__/Config.cpython-36.pyc
index e60788c1354940c10e4d9af28e7f2e0d9bee4291..670a87d685d254d5135178a568977be23b640b02 100644
GIT binary patch
delta 2776
zcmb7`Yj6`)6vulLy6FRGp#)4JO<Slm6iOeIM`>Y7J1{y7eWBJUgr#grASCf-MeBo7
zMU+Ri_5xb#2WBWDtuKf$2!f(KM*IK>zIG6`pA?l34vwS44Br1;D7zufXg>VroO91>
z?>&1r$9hS3FG(&PHZ0|86J2t2U-HI8%ZS40!NfTh;$eT-vZYB02?>7}vChPa(RoQr
zNGfzHP!R1)o@bS&z)%?dedu-b;$OpiEVGNSSPT$A;#lO)q(h<TfsjgoQlJ_r18SJF
z)+tp$yqDcrl*`ga-j(4d)D4TViMb;WTO_>>yD_Sp)Uj@R&4_8-&6BJu5p{7$l>@Yz
zU9~?r>vkd3&}_?Pwb|$O-4fJW0;Tna^;TNJ){m|#n2OS)fct@b-~o<!UHv*b(PCc;
zk0pSML$msWDqY4Thb!Ii09f|2GDjC%l{wV7z3j4M>zIf5a=iTQnJ&|Iv;MIcO~IST
z{WOjux4D#-6;V_#-xo&G#<1PatB!b1E;zCCwz2GV7Bi()T)RwqMs}}B|0Mf6liru}
znMq%lyF2}fn6_By_`W-hzm4B%ihOlKk4c}C*KN`V@&-)$*}DhQpNXw>eBuF<Q+s~Y
zq|bF<C{qyB0;~jlz%n2k$l-`{AZ)>;AB~x!r3Fq>`8*<C05$=efh`<vtC(X?s`7X^
zL{)E4UZc<stQt;>bxi(qTAU1;jq<S$QHACq)t#>dJ2a&`zlpDDX<?c2)T%<g0e?h|
z#B**fyl#)Ff>`=8+dMVX_(<&CsfUN{KuIg@g2#5YwkR)s8$@H<{Y9UW9u_J#2X8Ch
zH$QF}9`2FkNW{M~s0~rp%7L4kLy_1csY);sqNV~{N;+I`@Woq>uW?x;ev{9P((UPc
zx!q`tUMk&69D_}0wnl}Q4F&bX39$CEtXy$KFCn=_9FaDpd39JdWjS7Ug1AKPeiV2M
z7)(mvhPsLM&X_euY~wpz_arvTYZa=WKV#Qs<k?01e)t-<u5jj+HW4LbWjW9ctN_Fa
zF)n$MR4V&}-cVQ@5(q^i?j-$Ij1#0Kw<yXA?+Smjsu3?efINnRT-8$aWD!*fpW~M_
zdu7%x(!^$$=Sg}&)>1y%eh`5l#6(B;lwTvHn7v*xC3b=^nY5E#u5b|C-cIQ#4B}3=
zXFzNO#O*!;QY9;|c9Q9=soI%(9}EFt4P%v6#u1EuU%AQ3EBdP{vj!i-ZP6@+MH*RU
zd@^>TYNK=#Nlvl4dFu-Haq;IH{6|l7$vsoNJG2uH;xZRsq>)f#`;}rFl<aWhpM>h^
zX&x?mY%Kv<<#S3)*vaZ6?rzwui|ZOZ*7<b}bdONwN1N#>#0|Dh3!-~#CRo};KPiTO
zO1PVrC{rr8(lZEo7{CUKH{U0R70o8omHf3(#Xj9Ue~s`nrW#pI4<q11;0S<*jCHxH
zWrV4MsqjV@F8GDSZ-Smfpr|hT5NCKEwl4r2(O_vGy$IWv!2c<)%Km^?i5|S?5(&oL
z3B3t?2KCMO_a^AqaJukc+t|PGjHSF7-hAp4+zx#X^udKy6vA2P9f*O9nGzxfdST2!
zHz3{3SMqjC*dBhH`S{)K(@*3h<o*~i%mmv~yOiu;S86K?Cu2$c=MWQ$e;U*ec!3}g
z0v-V(thO$<?MaAF0Z#*Mz<Qt^cox_IJO^w9h67uH7Xk4?bV1q<^Z>m8?pyjM&<DH+
zybl}##9kdw;2qPm(7yt{0loz^;Pyibset9ww@N#(Qo?rDyK+9~q9;Ya?pp$($Edhn
hli-jF>|{UGkIFs9MNJAVU!^pw)B^h*c}F+o{{x^Ve82zz

delta 2705
zcma);ZBSHI7{@uguz)Bm<}R?guA-}~EboFZBa;YY6E(UBld*#96)v#2>@sIp9nvrj
zD9p+@nnTiPlTJleV|Ub1ld?B;s>#ycuG2K~shp<9scHIPQ|Es#!M)2`?T6ny&vTx)
zd(OGf-oc|}<S0olnKo@oy3yY_dLp?mft0b|6BkGW3FQVNvnz?SNv?D&v7e;m#>K@^
z1#YP~d8N^m4?}_U>D1fAtXbq%*~Ti;Y|^i(I}Bw-@G1rfVB%QqveBtf)C`bHfJHzR
zP{#gTZ8udzEN9=ZwzHe*_DuR9*Ik6VVAd3rY}3ygOlln_&ln^%EMlojFXe9D<cJuK
z)P*8KfL5`~mim;5LS-z^>PTD5z2wAALZqq9*Tsviw^rm9qQDGb9gqV&#1X6ZMQeh=
z;)BNqz`-FK{lN$o*o{m_)V|P-%$zll(HJ`yJCn6HqapU1)L?Vyj)3*ueNB5LOq}_H
zIgA#xJJV=SCzGtKd)8@P=nBV6I{jzIQJwz6?C*4X@0`mzy*~Sxd3#i8mO6KMvhmW~
zBf7}&yhA!Y%Q>ji&pStT`tbZw^JCGKh8CREIW^=+I^E{FTH1o3R-g^=0|GD`n8Oj{
za9Q@k8`?}#YObAFo<KxDum^Y&*vnS9Gc8X;)NDUj<qUSjoxytC*Uiy!ER>g}&CXcf
zAgO10`7?Fay8PSsM1|I)&H%C->6W0&pu7=s(pG``f<lKF-GY@}D6q0)1&5~{<S{NI
zWr%-{)f76-2Ow%%W_^Y4lOa}JWYaYfDmwANotp5PgmBp38kFP3twMmOSKdz~_L?GM
zFdU-V0^BTH!)#xHh3zV)j+eRqPskyx5x>@(O)oiM9_DtfQ950+mt>8%BPZ61yn#?q
zJ&z97P&zAHIhB1#KDImYO`AI-x-3Jb9}tI<`#1``4vZ(IZ$RC{b}n9;p=|e?T=yo_
z3R^|09{GOu*<z<fiI?E3-NfW2Hv>x4Cag38&436fw`M<=yh(gY_=CPsryL&$g~P5S
zb?agT$;mCE*z9ZeH$`OPqo<HZThO=k6u&ywWloctomDJzTTa94Y?Q0Cd)aS<6tiv1
z^A=7#trF#g>$4nke8kVMkum|5IXxXS*yO?b2qv?o8sc3*9Z(LWGh3ydRI(+NcJp!=
zFs^9_5M;sf3hjlG&X)HZc^6d`<+H}0!)1^S`GqoR*FD<diaygvNOF<ouI!%w78if~
zi~r~)F1eP_@vhZj9kXBcpG4DOiqfD3Y5B;nEK6I@MQ>tDK#1@wOH0_1%5$#gVKdfw
z)_RSrJZtGIiYjbrq8AZ2-ZphhJyr7zIZ8iMH2sutH#tF|ROq4~BLs7wZ{U?x)<xKs
zvGbM*KT|4`&Ga1B&jW7*sIMzyyY~+gh*^ey4qqj|QfUP07r<3OSFMeH3ENk|e^u<h
z?-#Q3Un=6?N`1CFSES#->DvFT$^B$=ClztmeN+qx(ogrFF_ab~9$yfyEcKu#&Rr>u
z%~p3oxp5QMfOv3YYhvBlTBI9$=_Z=Q?(oUgM#XTOdIs+z_XR+6=&WVU26B*nxMsO4
z4@=5NgYvzAPX*czYy?6;CqP+&Cwu3^5O)BN0Nuc&z%F1n&;#@WeZb>DD)1EW41o7^
zY>-150*(QgsPt9fHQ)qr5;y~l#1(PA2k{E<HSjGU1J{9xhg2jtyY1;R;je?g7u#zb
z>ptP4cZxdNTLPhN^f=a%fE3_G4%ry)3iGcozKcS$KiEp&fzcG&)Gjtfr~z)1%G|2S
F`5QP8T=M_`

diff --git a/__pycache__/distribute_training.cpython-36.pyc b/__pycache__/distribute_training.cpython-36.pyc
index 5e99d5273f2b6d8b5065515d755f38837c8a10a2..a3ff2fc1c5eb5bc9c870d4d689c948a060776797 100644
GIT binary patch
delta 348
zcmexgzO#bMn3tDJ+Oa5BF?}OfIWuF}<R<1edB!Z^8ip*9W=2K^AT650n$29a2gnkG
zuqJD>L`Sn_GfiMD5~yL!QYlddsmyZ%s!+;OZe~bf2g~Vz<<!7(YH&FYu$%{2P8}@A
z1!il2ZDR~((Bz!_kR@=kIqOAszP$L{%)HFWTx|Lx9L1?6@kyygc~P7oGCnOQF@3Wy
zn*t-F)#N;O874-a&3)|Wm>In`$8t+BGL}u|<1}<oNXsu$NG(guDNQWN%+FKE%}+_?
zQUHSFqSVBa)MADFyqrpfl8jWaOpZcgUW!6`PJUuiVva&_Nov95<va(ta~K#Hilu-A
w1MB7@-mgrIZJRR$R?0D6-z=o>!N@4KIoe<XJLAI5Z_IgE*tIy6IHWkF0d-YkA^-pY

delta 271
zcmdm4@xPqQn3tDJJ1#T!b=pR*a%RS;$xX~{$_!b;MU!e6vqVx@n;97yY8bLaQ`oYZ
ziw*!;Vi4A3b(Uy*_H3pJj71_ql`18wAeDJeKov?^%FPTZ9AG&Eu$&rLP7N-{36}Gj
ze2_(&k!$i@mOvI>1_p-7x~!L&IP)g6vFVF&CZ!hT#i!*YrblrUr<TNnxSKuM6c`z8
zCTFqBurRSPa&7KlKgZ1IvpJkwf|1c|asiLwWJ^xI$?thKapwVzE0zKh3@n@1@P1`t
pY~Q>{V5J=6jm<v#9*m6Qo2ME~U}s#kS<Qlng<Xe3i9?b@8UV$cNUZ<>

diff --git a/base/Base.cpp b/base/Base.cpp
index 3417a7f..8d5030c 100644
--- a/base/Base.cpp
+++ b/base/Base.cpp
@@ -68,7 +68,9 @@ struct Parameter {
 	INT negRelRate;
 };
 
-
+/*
+ * feed batch with training triples / corrupted triples
+*/
 void* getBatch(void* con) {
 	Parameter *para = (Parameter *)(con);
 	INT id = para -> id;
@@ -94,7 +96,7 @@ void* getBatch(void* con) {
 		INT i;
 
 		/**
-		* select batch triple / train triple
+		* select new batch triple / train triple
 		**/
 		if (newBatchTotal > 0){
             i = rand_max_range(id, trainTotal_ - newBatchTotal, trainTotal_);
@@ -137,12 +139,13 @@ void* getBatch(void* con) {
 		}
 	}
 
-
 	pthread_exit(NULL);
 }
 
 
-
+/*
+ * Sample a batch to use during training
+*/
 extern "C"
 void sampling(INT *batch_h, INT *batch_t, INT *batch_r, REAL *batch_y, INT batchSize, INT negRate = 1, INT negRelRate = 0) {
 	pthread_t *pt = (pthread_t *)malloc(workThreads * sizeof(pthread_t));
diff --git a/base/Reader.h b/base/Reader.h
index 25966d6..d853670 100644
--- a/base/Reader.h
+++ b/base/Reader.h
@@ -17,8 +17,6 @@ Triple *trainHead;
 Triple *trainTail;
 Triple *trainRel;
 
-
-//EDIT
 Triple *trainList_no;   //data structure containing training triples not ordered
 
 
@@ -30,7 +28,7 @@ void importTrainFiles() {
 
 	printf("The toolkit is importing datasets.\n");
 	FILE *fin;
-	FILE *fin_; //EDIT
+	FILE *fin_;
 	int tmp;
 
     //read relations
@@ -58,7 +56,7 @@ void importTrainFiles() {
 	/**
     * EDIT
     * check if new batch file exists
-    * if it exists get only new batch size
+    * if it exists get only new batch size (another script has already incorporeted the new triples into train2id.txt)
     **/
     fin_ = fopen((inPath + "batch2id.txt").c_str(), "r");
     if (fin_ != nullptr){
@@ -76,7 +74,7 @@ void importTrainFiles() {
 		return;
 	}
 	tmp = fscanf(fin, "%ld", &trainTotal);
-	trainTotal_ = trainTotal;           //traingTotal_ contains # training triples counting duplicates, too
+	trainTotal_ = trainTotal;           //traingTotal_ contains # training triples counting duplicates too
 	printf("The total of train triples is %ld.\n", trainTotal);
 
 
@@ -87,7 +85,7 @@ void importTrainFiles() {
 	trainRel = (Triple *)calloc(trainTotal, sizeof(Triple));
 	freqRel = (INT *)calloc(relationTotal, sizeof(INT));
 	freqEnt = (INT *)calloc(entityTotal, sizeof(INT));
-	//EDIT
+	//not ordered
     trainList_no = (Triple *)calloc(trainTotal, sizeof(Triple));    //not ordered
 
 	for (INT i = 0; i < trainTotal; i++) {
@@ -95,7 +93,7 @@ void importTrainFiles() {
 		tmp = fscanf(fin, "%ld", &trainList[i].t);
 		tmp = fscanf(fin, "%ld", &trainList[i].r);
 
-		//EDIT
+
         trainList_no[i].h = trainList[i].h;
         trainList_no[i].t = trainList[i].t;
         trainList_no[i].r = trainList[i].r;
diff --git a/base/Setting.h b/base/Setting.h
index 2df9cf5..7202f44 100644
--- a/base/Setting.h
+++ b/base/Setting.h
@@ -114,7 +114,8 @@ void setBern(INT con) {
 /*
 ============================================================
 */
-REAL interval = 0.01;
+
+REAL interval = 0.01;   //used to compute thresholds during triple classification task
 
 
 #endif
diff --git a/commands.txt b/colab/commands.txt
similarity index 72%
rename from commands.txt
rename to colab/commands.txt
index 8c23688..be3487b 100644
--- a/commands.txt
+++ b/colab/commands.txt
@@ -1,4 +1,6 @@
 #commands to run OpenKEonSpark on Google Colaboratory platform
+#before running, create a folder in your Google Drive named DBpedia
+#uplod in the folder the zip archive containing the dataset given as output from generate.py
 
 #install
 !apt-get install openjdk-8-jdk-headless -qq > /dev/null
@@ -15,12 +17,10 @@ drive.mount('/content/drive')
 
 #clone project
 !git clone https://github.com/luigiba/OpenKEonSpark.git
-!mv OpenKEonSpark/gpu_info.py /usr/local/lib/python3.6/dist-packages/tensorflowonspark/gpu_info.py
+!mv OpenKEonSpark/colab/gpu_info.py /usr/local/lib/python3.6/dist-packages/tensorflowonspark/gpu_info.py
 
 #unzip dataset
-!unzip /content/drive/My\ Drive/DBpedia/5.zip -d /content/drive/My\ Drive/DBpedia/
 !unzip /content/drive/My\ Drive/DBpedia/10.zip -d /content/drive/My\ Drive/DBpedia/
-!unzip /content/drive/My\ Drive/DBpedia/15.zip -d /content/drive/My\ Drive/DBpedia/
 
 #set env vars
 import os
@@ -35,9 +35,7 @@ os.environ["WORK_DIR_PREFIX"] = "/content/OpenKEonSpark"
 os.environ["SPARK_HOME"] = "/content/spark-2.1.1-bin-hadoop2.7"
 
 #execute
-!bash $WORK_DIR_PREFIX/run_dbpedia.sh 5 64 "TransE" 0.0001
-!bash $WORK_DIR_PREFIX/run_dbpedia.sh 10 64 "TransE" 0.0001
-!bash $WORK_DIR_PREFIX/run_dbpedia.sh 15 64 "TransE" 0.0001
+!bash $WORK_DIR_PREFIX/colab/run_dbpedia.sh 10 64 "TransE" 0.00001
 
 
 
diff --git a/gpu_info.py b/colab/gpu_info.py
similarity index 100%
rename from gpu_info.py
rename to colab/gpu_info.py
diff --git a/colab/run_dbpedia.sh b/colab/run_dbpedia.sh
new file mode 100644
index 0000000..c1778f6
--- /dev/null
+++ b/colab/run_dbpedia.sh
@@ -0,0 +1,125 @@
+#this is an example of script that can be used on google colab to train the embedding and evaluate them
+# the link prediction evaluation will be performed only for the last batch
+# for the other batch it will be performed only triple classification evaluation
+#before starting the script:
+#   run split.py to set the dataset in a properly format
+#   set environment variables, e.g.:
+#     $JAVA_HOME = "/usr/lib/jvm/java-8-openjdk-amd64"
+#     $SPARK_WORKER_INSTANCES = '3'
+#     $PYSPARK_PYTHON = '/usr/bin/python3'
+#     $CUDA_VISIBLE_DEVICES = "0"
+#     $CORES_PER_WORKER = "1"
+#     $MEMORY_PER_WORKER = "4g"
+#     $LIB_CUDA = "/usr/local/cuda-10.0/lib64"
+#     $WORK_DIR_PREFIX = "/content/OpenKEonSpark"
+#     $SPARK_HOME = "/content/spark-2.1.1-bin-hadoop2.7"
+
+
+echo "====================================== Parameters ======================================"
+echo "Number of batches: $1"
+echo "Embedding dimensionality: $2"
+echo "Model to use: $3"
+echo "Learning rate: $4"
+
+echo "====================================== Clearning res_spark directory ======================================"
+rm $WORK_DIR_PREFIX/res_spark/*
+
+echo "====================================== Stopping Spark Master & slaves ======================================"
+$SPARK_HOME/sbin/stop-slave.sh
+$SPARK_HOME/sbin/stop-master.sh
+
+
+echo "====================================== Starting Spark Master & slaves ======================================"
+$SPARK_HOME/sbin/start-master.sh; $SPARK_HOME/sbin/start-slave.sh -c $CORES_PER_WORKER -m $MEMORY_PER_WORKER spark://$(hostname):7077
+n=$1
+m=$((n-1))
+
+#iterate over batches
+for i in `seq 0 $m`
+do
+  #restore link prediction evaluation if evaluation checkpoints have been founded
+  if [ -f /content/drive/My\ Drive/DBpedia/$n/$i/model/thread0 ]; then
+    echo "====================================== Test for batch $i ======================================"
+	  $SPARK_HOME/bin/spark-submit --master spark://$(hostname):7077 \
+    --py-files $WORK_DIR_PREFIX/distribute_training.py,$WORK_DIR_PREFIX/Config.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransE.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransH.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransR.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransD.py \
+    --driver-library-path=$LIB_CUDA --conf spark.dynamicAllocation.enabled=false --conf spark.task.cpus=$CORES_PER_WORKER --executor-memory $MEMORY_PER_WORKER --num-executors $SPARK_WORKER_INSTANCES \
+	  $WORK_DIR_PREFIX/main_spark.py \
+    --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 --cpp_lib_path $WORK_DIR_PREFIX/release/Base.so \
+    --embedding_dimension $2 --model $3 --mode evaluation | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
+
+	  continue
+  fi
+
+
+  if [ -f /content/drive/My\ Drive/DBpedia/$n/$i/res.txt ]; then
+    echo "Batch $i already done; Skipping batch $i"
+	  continue
+  fi
+
+
+  if [ -f /content/drive/My\ Drive/DBpedia/$n/$i/model/checkpoint ]; then
+    echo "====================================== Test for batch $i ======================================"
+    if [ $i -eq $m ]; then
+	    $SPARK_HOME/bin/spark-submit --master spark://$(hostname):7077 \
+    --py-files $WORK_DIR_PREFIX/distribute_training.py,$WORK_DIR_PREFIX/Config.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransE.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransH.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransR.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransD.py \
+    --driver-library-path=$LIB_CUDA --conf spark.dynamicAllocation.enabled=false --conf spark.task.cpus=$CORES_PER_WORKER --executor-memory $MEMORY_PER_WORKER --num-executors $SPARK_WORKER_INSTANCES \
+	  $WORK_DIR_PREFIX/main_spark.py \
+    --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 --cpp_lib_path $WORK_DIR_PREFIX/release/Base.so \
+    --embedding_dimension $2 --model $3 --mode evaluation | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
+    else
+      python3 $WORK_DIR_PREFIX/test.py /content/drive/My\ Drive/DBpedia/$n/$i/ /content/drive/My\ Drive/DBpedia/$n/$i/model/ $WORK_DIR_PREFIX/release/Base.so $2 $3 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
+    fi
+  fi
+
+
+  if [ $i != 0 ]; then
+    k=$((i-1))
+    cd $WORK_DIR_PREFIX/res_spark
+    if [ "$(ls -1A | wc -l)" -eq 0 ] ; then
+      echo "Copying model into res_spark dir"
+      cp /content/drive/My\ Drive/DBpedia/$n/$k/model/* $WORK_DIR_PREFIX/res_spark/
+    fi
+
+    cd /content/drive/My\ Drive/DBpedia/$n/$i
+    if [ "$(ls -1A | wc -l)" -le 10 ] ; then
+      echo "Copying data into new batch dir"
+		  cp /content/drive/My\ Drive/DBpedia/$n/$k/entity2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/relation2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/test2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/valid2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/train2id.txt /content/drive/My\ Drive/DBpedia/$n/$i/
+    fi
+
+    cd /content
+  fi
+
+
+	echo "====================================== Starting Training for batch $i ======================================"
+	$SPARK_HOME/bin/spark-submit --master spark://$(hostname):7077 \
+	--py-files $WORK_DIR_PREFIX/distribute_training.py,$WORK_DIR_PREFIX/Config.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransE.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransH.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransR.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransD.py \
+    --driver-library-path=$LIB_CUDA --conf spark.dynamicAllocation.enabled=false --conf spark.task.cpus=$CORES_PER_WORKER --executor-memory $MEMORY_PER_WORKER \
+    --num-executors $SPARK_WORKER_INSTANCES \
+	$WORK_DIR_PREFIX/main_spark.py \
+    --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
+
+
+	echo "====================================== Copying model for batch $i ======================================"
+	cp $WORK_DIR_PREFIX/res_spark/* /content/drive/My\ Drive/DBpedia/$n/$i/model/
+
+	
+	echo "====================================== Test for batch $i ======================================"
+	if [ $i -eq $m ]; then
+	  $SPARK_HOME/bin/spark-submit --master spark://$(hostname):7077 \
+    --py-files $WORK_DIR_PREFIX/distribute_training.py,$WORK_DIR_PREFIX/Config.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransE.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransH.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransR.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransD.py \
+    --driver-library-path=$LIB_CUDA --conf spark.dynamicAllocation.enabled=false --conf spark.task.cpus=$CORES_PER_WORKER --executor-memory $MEMORY_PER_WORKER --num-executors $SPARK_WORKER_INSTANCES \
+	  $WORK_DIR_PREFIX/main_spark.py \
+    --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 --cpp_lib_path $WORK_DIR_PREFIX/release/Base.so \
+    --embedding_dimension $2 --model $3 --mode evaluation | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
+	else
+	  python3 $WORK_DIR_PREFIX/test.py /content/drive/My\ Drive/DBpedia/$n/$i/ /content/drive/My\ Drive/DBpedia/$n/$i/model/ $WORK_DIR_PREFIX/release/Base.so $2 $3 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
+	fi
+
+done
+
diff --git a/distribute_training.py b/distribute_training.py
index b352f10..c1ccc5b 100644
--- a/distribute_training.py
+++ b/distribute_training.py
@@ -43,12 +43,10 @@ def get_conf(argv=None):
     else: con.set_test_link_prediction(True)
 
     con.set_train_times(argv.train_times)
-    con.set_nbatches(argv.n_batches)
+    con.set_nbatches(argv.n_mini_batches)
     con.set_alpha(argv.alpha)
     con.set_margin(argv.margin)
-
-    if argv.bern_flag:
-        con.set_bern(1)
+    con.set_bern(argv.bern_flag)
 
     if argv.ent_dimension != 0 and argv.rel_dimension != 0:
         con.set_ent_dimension(argv.ent_dimension)
@@ -109,7 +107,8 @@ def create_model(con):
 
 def create_model_evaluation(con):
     '''
-    create the model using the Config parameters
+    create the model using the Config parameters for evaluation mode
+    creates only the model and gloabal step
     '''
     with tf.variable_scope("", reuse=None, initializer = tf.contrib.layers.xavier_initializer(uniform = True)):
         trainModel = con.model(config = con)
diff --git a/main_spark.py b/main_spark.py
index 707ec36..084e3fe 100644
--- a/main_spark.py
+++ b/main_spark.py
@@ -21,7 +21,6 @@
 TEST_FILE_NAME = 'test2id.txt'
 VALID_FILE_NAME = 'valid2id.txt'
 
-#TODO: retrieve these from models methods
 #tensors which depends from entities dimension
 ENTITY_EMBEDDING_TENSOR_NAME = 'ent_embeddings'     #from TransE, TransH, TransR, TransD
 ENTITY_TRANSFER_TENSOR_NAME = 'ent_transfer'        #from TransD
@@ -31,7 +30,6 @@ def update_entities_and_model():
     '''
     Update the tensor variables if new entites are introduced in the new batch
     '''
-
     n_entities = 0
     n_new_entities = 0
     final_entity_size = 0
@@ -209,10 +207,11 @@ def get_test_total():
 
 
 def n_n():
-    if sys.argv.debug: print("Generating constraints file...")
     '''
-    Generate type constrain file
+    Generates type_constrain.txt file
     '''
+    if sys.argv.debug: print("Generating constraints file...")
+
     lef = {}
     rig = {}
     rellef = {}
@@ -301,36 +300,36 @@ def n_n():
     parser.add_argument("--cluster_size", help="number of nodes in the cluster", type=int, default=int(sc._conf.get("spark.executor.instances")))
     parser.add_argument("--num_ps", help="number of ps nodes", type=int, default=1)
     parser.add_argument("--num_gpus", help="number of gpus to use", type=int, default=0)
-    parser.add_argument("--cpp_lib_path", help="cpp lib.so absolute path", type=str, default=None)
-    parser.add_argument("--input_path", help="dataset absolute path", type=str, default="drive/My\ Drive/DBpedia/5/0/")
-    parser.add_argument("--output_path", help="model output absolute path", type=str, default="OpenKE_new_Spark/res_spark")
-    parser.add_argument("--train_times", help="no. epochs", type=int, default=1000)
-    parser.add_argument("--n_batches", help="no. batches", type=int, default=0)
-    parser.add_argument("--alpha", help="learning rate", type=float, default=0.001)
+    parser.add_argument("--cpp_lib_path", help="cpp lib.so absolute path contained in ./release/Base.so", type=str, default=None)
+    parser.add_argument("--input_path", help="dataset absolute path", type=str, default=None)
+    parser.add_argument("--output_path", help="model output absolute path", type=str, default=None)
+    parser.add_argument("--train_times", help="no. epochs", type=int, default=100)
+    parser.add_argument("--n_mini_batches", help="# mini batches; if set to zero it will be automatically computed", type=int, default=0)
+    parser.add_argument("--alpha", help="learning rate", type=float, default=0.00001)
     parser.add_argument("--margin", help="margin hyperparameter used during training", type=float, default=1.0)
-    parser.add_argument("--bern_flag", help="use bern method for sampling", type=bool, default=False)
+    parser.add_argument("--bern_flag", help="whether to use or not bern method for sampling; ; 0=False, n=True", type=int, default=0)
     parser.add_argument("--embedding_dimension", help="embedding dimension (both entities and rel)", type=int, default=64)
     parser.add_argument("--ent_dimension", help="entities embedding dimension", type=int, default=0)
     parser.add_argument("--rel_dimension", help="relations embedding dimension", type=int, default=0)
     parser.add_argument("--ent_neg_rate", help="number of negative triples generated by corrupting the entity", type=int, default=1)
     parser.add_argument("--rel_neg_rate", help="number of negative triples generated by corrupting the realtion", type=int, default=0)
     parser.add_argument("--optimizer", help="Optimization algorithm", type=str, default="SGD")
-    parser.add_argument("--early_stop_patience", help="no. epochs to wait for loss improvement before early stop", type=int, default=5)
+    parser.add_argument("--early_stop_patience", help="no. epochs to wait for accuracy/loss improvement before early stop", type=int, default=5)
     parser.add_argument("--early_stop_stopping_step", help="perfrom early stop each stopping step", type=int, default=1)
     parser.add_argument("--early_stop_start_step", help="perfrom early stop from start step", type=int, default=1)
     parser.add_argument("--model", help="model to be used", type=str, default="TransE")
     parser.add_argument("--debug", help="if Ture prints additional debug information", type=bool, default=True)
-    parser.add_argument("--mode", help="train/evaluation", type=str, default="train")
+    parser.add_argument("--mode", help="whether to perform train or evaluation mode", type=str, default="train")
     parser.add_argument("--test_head", help="perform link prediction evaluation on missing head, too (only if mode != 'train'); 0=False, n=True", type=int, default=0)
     (args, remainder) = parser.parse_known_args()
     num_workers = args.cluster_size - args.num_ps
     print("===== num_executors={}, num_workers={}, num_ps={}".format(args.cluster_size, num_workers, args.num_ps))
     sys.argv = args
 
-
+    #generate type_constrain.txt file
+    n_n()
     if args.mode == 'train':
         if is_new_batch(): feed_batch()
-        n_n()
         try: os.remove(os.path.join(args.output_path,"stop.txt"))
         except: pass
 
diff --git a/run_dbpedia.sh b/run_dbpedia.sh
deleted file mode 100644
index 4129e27..0000000
--- a/run_dbpedia.sh
+++ /dev/null
@@ -1,87 +0,0 @@
-echo "====================================== Params ======================================"
-echo "$1"
-echo "$2"
-echo "$3"
-echo "$4"
-
-echo "====================================== Clearning res_spark directory ======================================"
-rm $WORK_DIR_PREFIX/res_spark/*
-
-echo "====================================== Stopping Spark Master & slaves ======================================"
-$SPARK_HOME/sbin/stop-slave.sh
-$SPARK_HOME/sbin/stop-master.sh
-
-
-echo "====================================== Starting Spark Master & slaves ======================================"
-$SPARK_HOME/sbin/start-master.sh; $SPARK_HOME/sbin/start-slave.sh -c $CORES_PER_WORKER -m $MEMORY_PER_WORKER spark://$(hostname):7077
-n=$1
-m=$((n-1))
-
-for i in `seq 0 $m`
-do
-  if [ -f /content/drive/My\ Drive/DBpedia/$n/$i/model/thread0 ]; then
-    echo "====================================== Test for batch $i ======================================"
-	  python3 $WORK_DIR_PREFIX/test.py $i $n $2 $3 1 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
-	  continue
-  fi
-
-  if [ -f /content/drive/My\ Drive/DBpedia/$n/$i/res.txt ]; then
-    echo "Batch $i already done; Skipping batch $i"
-	  continue
-  fi
-
-  if [ -f /content/drive/My\ Drive/DBpedia/$n/$i/model/checkpoint ]; then
-    echo "====================================== Test for batch $i ======================================"
-    if [ $i -eq $m ]; then
-	    python3 $WORK_DIR_PREFIX/test.py $i $n $2 $3 1 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
-    else
-      python3 $WORK_DIR_PREFIX/test.py $i $n $2 $3 0 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
-    fi
-  fi
-
-
-  if [ $i != 0 ]; then
-    k=$((i-1))
-    cd $WORK_DIR_PREFIX/res_spark
-    if [ "$(ls -1A | wc -l)" -eq 0 ] ; then
-      echo "Copying model into res_spark dir"
-      cp /content/drive/My\ Drive/DBpedia/$n/$k/model/* $WORK_DIR_PREFIX/res_spark/
-    fi
-
-    cd /content/drive/My\ Drive/DBpedia/$n/$i
-    if [ "$(ls -1A | wc -l)" -le 10 ] ; then
-      echo "Copying data into new batch dir"
-		  cp /content/drive/My\ Drive/DBpedia/$n/$k/entity2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/relation2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/test2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/valid2id.txt /content/drive/My\ Drive/DBpedia/$n/$k/train2id.txt /content/drive/My\ Drive/DBpedia/$n/$i/
-    fi
-
-    cd /content
-  fi
-
-
-	echo "====================================== Starting Training for batch $i ======================================"
-	$SPARK_HOME/bin/spark-submit --master spark://$(hostname):7077 \
-	--py-files $WORK_DIR_PREFIX/distribute_training.py,$WORK_DIR_PREFIX/Config.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransE.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransH.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransR.py,$WORK_DIR_PREFIX/Model.py,$WORK_DIR_PREFIX/TransD.py \
-    --driver-library-path=$LIB_CUDA \
-    --conf spark.dynamicAllocation.enabled=false \
-    --conf spark.task.cpus=$CORES_PER_WORKER --executor-memory $MEMORY_PER_WORKER \
-    --num-executors $SPARK_WORKER_INSTANCES \
-	$WORK_DIR_PREFIX/main_spark.py \
-    --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 --n_batches 100
-
-
-	echo "====================================== Copying model for batch $i ======================================"
-	cp $WORK_DIR_PREFIX/res_spark/* /content/drive/My\ Drive/DBpedia/$n/$i/model/
-
-	
-	echo "====================================== Test for batch $i ======================================"
-	if [ $i -eq $m ]; then
-	  python3 $WORK_DIR_PREFIX/test.py $i $n $2 $3 1 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
-	else
-	  python3 $WORK_DIR_PREFIX/test.py $i $n $2 $3 0 | tee /content/drive/My\ Drive/DBpedia/$n/$i/res.txt
-	fi
-
-done
-
diff --git a/split/generate.py b/split/generate.py
index 3400942..d526237 100644
--- a/split/generate.py
+++ b/split/generate.py
@@ -3,45 +3,61 @@
 """
 Created on Sat Jun 29 20:10:44 2019
 
-You can use this script to split your dataset into batches 
-to feed into OpenKEonSpark
+
+You can use this script to
+    assign numerical identifiers to resources
+    split data into training/test/validation set
+    split data into batches
+    generate ontology_constrain file
+This script will output a directory which is named as n, where n is number of batches used to split the dataset.
+Inside the directory there will be n-1 (named from 0 to n-1) sub-folders. Each sub-folder contains a different batch.
+Moreover, inside each sub-batch folder i, it will be created the folder "model" that can be used to save the model up to batch i.
 
 @author: Luigi Baldari
+
+@:param
+    TARGET_RELATION
+        Specify in this array the target relation/s
+        Training set will contain all the relations (target included)
+        test set and validation set will contain ontly target relation/s
+    path_t
+        Path to N-Triples dataset containing all the triples with only the target relation/s
+    path_r
+        Path to N-Triples dataset containing all the other triples (except the target relation/s)
+    path_h
+        Path which contains the class hierarchy (from which the ontology_constrain.txt will be created)
+        The root node (i.e. the first line of the file) contains the most general class. As
+        new specific classes are met, a new level of indentation has to be added to the new line.
+        In this way the classes which have the same number of tabs are on the same level of the tree.
+    SKIP_DATA_PROPERTY
+        if set to True, triples with Data Propertiy relations will be discarded and only triples with Object Propertiy relations will be taken into account
+    GENERATE_ONTOLOGY_FILES
+        if set to True, it will generate ontology_constrain.txt
+    N_BATCHES
+        number of batches to split the dataset into
+    TEST_SET_PERCENTAGE
+        percentage of triples with target relation/s for each batch test set
+    VALIDATION_SET_PERCENTAGE
+        percentage of triples with target relation/s for each batch validation set
+
+
+
 """
+
 import math
 from random import shuffle
 import os
 
-#The relation of interest
-TARGET_RELATION = '<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>'
 
-#dataset containing all the triples with only the target relation
+TARGET_RELATION = ['<http://www.w3.org/1999/02/22-rdf-syntax-ns#type>']
 path_t = "/home/luigi/Desktop/DBpedia/instancetype_en.nt"
-
-#dataset containing all the other triples (except the target relation)
 path_r = "/home/luigi/Desktop/DBpedia/newstrictmappingbased_en.nt"
-
-#path which contains the class hierarchy
 path_h = "class_hierarchy.txt"
-
-#path that will contains the ontology constraints for each batch
-path_ont_con = "ontology_constrain.txt"
-
-#if set to True, triples with Data Propertiy relations will be discarded
-# only triples with Object Propertiy relations will be taken into account
 SKIP_DATA_PROPERTY = True
-
-#if set to True, it will generate ontology files constraints
 GENERATE_ONTOLOGY_FILES = True
-
-#number of batches to split the dataset into
-N_BATCHES = 1
-
-#percentage of triples with target relation for each batch test set 
+N_BATCHES = 10
 TEST_SET_PERCENTAGE = 10
-
-#percentage of triples with target relation for each batch validation set 
-VALIDATION_SET_PERCENTAGE = 10 
+VALIDATION_SET_PERCENTAGE = 10
 
 
 
@@ -187,7 +203,7 @@
     try: os.mkdir(str(N_BATCHES)+'/'+str(b))
     except: pass
     
-    try: os.mkdir((str(N_BATCHES)+'/'+str(b)+'/model'))
+    try: os.mkdir((str(N_BATCHES)+'/'+str(b)+'/model'))     #create dir where the model will be stored
     except: pass
     
     if b+1 == N_BATCHES: rig = n_triples
@@ -231,7 +247,7 @@
     
     if GENERATE_ONTOLOGY_FILES:
         print(" LOG:\tGenerating ontology constrain file")
-        with open("./"+str(N_BATCHES)+'/'+str(b)+"/"+path_ont_con, 'w') as f_ont_con:
+        with open("./"+str(N_BATCHES)+'/'+str(b)+"/ontology_constrain.txt", 'w') as f_ont_con:
             classes_len = 0
             for key in classes.keys():
                 if entities[key] <= entities[ent]: 
@@ -272,7 +288,7 @@
     test_triples_ids = []
     remaining = []
     for i in batch_triples_ids:
-        if triples[i].split(" ")[1].strip() == TARGET_RELATION:
+        if triples[i].split(" ")[1].strip() in TARGET_RELATION:
             test_triples_ids.append(i)
         else:
             remaining.append(i)
diff --git a/test.py b/test.py
index 439b8ca..6f896dc 100644
--- a/test.py
+++ b/test.py
@@ -1,3 +1,19 @@
+"""
+Once learned the embedding model, this script could be used to:
+    use the model learnt (predict_head_entity, predict_tail_entity, predict_relation, predict_triple)
+    evaluate test triple classification task:
+        - precision, recall, accuracy, f-measure
+        - ROC curve
+
+@:param
+    sys.argv[1]; path where dataset is located
+    sys.argv[2]; path where the model is located
+    sys.argv[3]; path to OpenKEonSpark/release/Base.so
+    sys.argv[4]: embedding dimensionality
+    sys.argv[5]: model to use (transe/transh/transr/transd)
+    sys.argv[6] (optional): target relation index (if you want to plot ROC curve)
+"""
+
 from Config import Config
 from TransE import TransE
 from TransH import TransH
@@ -5,15 +21,17 @@
 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
 
 dataset_path = sys.argv[1]
 model_path = sys.argv[2]
 cpp_path = sys.argv[3]
-dim = sys.argv[4]
-model = sys.argv[5]
+dim = sys.argv[3]
+model = sys.argv[4]
 if (len(sys.argv) >= 7): target_rel_index = sys.argv[6]