From 08852b31280487c30fbac13e6f878a6a13c09442 Mon Sep 17 00:00:00 2001
From: t170815518 <christang_1023@outlook.com>
Date: Wed, 30 Sep 2020 14:18:08 +0800
Subject: [PATCH] Fix bug when entities include whitespace, by improving file
 reading process with ifstream

---
 TransE/Train_TransE.cpp | 141 +++++++++++++++++++++++-----------------
 1 file changed, 80 insertions(+), 61 deletions(-)

diff --git a/TransE/Train_TransE.cpp b/TransE/Train_TransE.cpp
index 65fe39d..f510969 100644
--- a/TransE/Train_TransE.cpp
+++ b/TransE/Train_TransE.cpp
@@ -8,6 +8,7 @@
 #include<ctime>
 #include<cmath>
 #include<cstdlib>
+#include <fstream>
 using namespace std;
 
 
@@ -73,8 +74,8 @@ class Train{
     void run(int n_in,double rate_in,double margin_in,int method_in)
     {
         n = n_in;
-        rate = rate_in;
-        margin = margin_in;
+        rate = rate_in;
+        margin = margin_in;
         method = method_in;
         relation_vec.resize(relation_num);
 		for (int i=0; i<relation_vec.size(); i++)
@@ -147,9 +148,9 @@ class Train{
              		for (int k=0; k<batchsize; k++)
              		{
 						int i=rand_max(fb_h.size());
-						int j=rand_max(entity_num);
-						double pr = 1000*right_num[fb_r[i]]/(right_num[fb_r[i]]+left_num[fb_r[i]]);
-						if (method ==0)
+						int j=rand_max(entity_num);
+						double pr = 1000*right_num[fb_r[i]]/(right_num[fb_r[i]]+left_num[fb_r[i]]);
+						if (method ==0)
                             pr = 500;
 						if (rand()%1000<pr)
 						{
@@ -242,47 +243,65 @@ class Train{
 Train train;
 void prepare()
 {
-    FILE* f1 = fopen("../data/entity2id.txt","r");
-	FILE* f2 = fopen("../data/relation2id.txt","r");
-	int x;
-	while (fscanf(f1,"%s%d",buf,&x)==2)
+    ifstream f1;
+    ifstream f2;
+    int x;
+    string line;
+
+	f1.open("../data/entity2id.txt");
+	while (getline(f1, line))
 	{
-		string st=buf;
-		entity2id[st]=x;
-		id2entity[x]=st;
-		entity_num++;
+        string st = line.substr(0, line.find('\t'));
+        x = stoi(line.substr(line.find('\t') + 1, line.length() - line.find('\t')));
+        entity2id[st] = x;
+        id2entity[x] = st;
+        entity_num++;
 	}
-	while (fscanf(f2,"%s%d",buf,&x)==2)
+	f1.close();
+
+	f2.open("../data/relation2id.txt");
+	while (getline(f2, line))
 	{
-		string st=buf;
-		relation2id[st]=x;
-		id2relation[x]=st;
-		relation_num++;
+        string st = line.substr(0, line.find('\t'));
+        x = stoi(line.substr(line.find('\t') + 1, line.length() - line.find('\t')));
+        relation2id[st] = x;
+        id2relation[x] = st;
+        relation_num++;
 	}
-    FILE* f_kb = fopen("../data/train.txt","r");
-	while (fscanf(f_kb,"%s",buf)==1)
+
+    ifstream f_kb;
+	f_kb.open("../data/train.txt");
+    string* s = new string[3];
+	while (getline(f_kb, line))
     {
-        string s1=buf;
-        fscanf(f_kb,"%s",buf);
-        string s2=buf;
-        fscanf(f_kb,"%s",buf);
-        string s3=buf;
-        if (entity2id.count(s1)==0)
+        // parse e1, rel, e2
+        size_t pos = 0;
+        int i = 0;
+        string token;
+        string delimiter = "\t";
+        while ((pos = line.find(delimiter)) != std::string::npos) {
+            token = line.substr(0, pos);
+            s[i++] = token;
+            line.erase(0, pos + delimiter.length());
+        }
+        s[2] = line;
+
+        if (entity2id.count(s[0])==0)
         {
-            cout<<"miss entity:"<<s1<<endl;
+            cout<<"miss entity:"<<s[0]<<endl;
         }
-        if (entity2id.count(s2)==0)
+        if (entity2id.count(s[2])==0)
         {
-            cout<<"miss entity:"<<s2<<endl;
+            cout<<"miss entity:"<<s[2]<<endl;
         }
-        if (relation2id.count(s3)==0)
+        if (relation2id.count(s[1])==0)
         {
-            relation2id[s3] = relation_num;
+            relation2id[s[1]] = relation_num;
             relation_num++;
         }
-        left_entity[relation2id[s3]][entity2id[s1]]++;
-        right_entity[relation2id[s3]][entity2id[s2]]++;
-        train.add(entity2id[s1],entity2id[s2],relation2id[s3]);
+        left_entity[relation2id[s[1]]][entity2id[s[0]]]++;
+        right_entity[relation2id[s[1]]][entity2id[s[2]]]++;
+        train.add(entity2id[s[0]],entity2id[s[2]],relation2id[s[1]]);
     }
     for (int i=0; i<relation_num; i++)
     {
@@ -307,38 +326,38 @@ void prepare()
     cout<<"relation_num="<<relation_num<<endl;
     cout<<"entity_num="<<entity_num<<endl;
     fclose(f_kb);
-}
-
-int ArgPos(char *str, int argc, char **argv) {
-  int a;
-  for (a = 1; a < argc; a++) if (!strcmp(str, argv[a])) {
-    if (a == argc - 1) {
-      printf("Argument missing for %s\n", str);
-      exit(1);
-    }
-    return a;
-  }
-  return -1;
+}
+
+int ArgPos(char *str, int argc, char **argv) {
+  int a;
+  for (a = 1; a < argc; a++) if (!strcmp(str, argv[a])) {
+    if (a == argc - 1) {
+      printf("Argument missing for %s\n", str);
+      exit(1);
+    }
+    return a;
+  }
+  return -1;
 }
 
 int main(int argc,char**argv)
 {
     srand((unsigned) time(NULL));
-    int method = 1;
-    int n = 100;
-    double rate = 0.001;
-    double margin = 1;
-    int i;
-    if ((i = ArgPos((char *)"-size", argc, argv)) > 0) n = atoi(argv[i + 1]);
-    if ((i = ArgPos((char *)"-margin", argc, argv)) > 0) margin = atoi(argv[i + 1]);
-    if ((i = ArgPos((char *)"-method", argc, argv)) > 0) method = atoi(argv[i + 1]);
-    cout<<"size = "<<n<<endl;
-    cout<<"learing rate = "<<rate<<endl;
-    cout<<"margin = "<<margin<<endl;
-    if (method)
-        version = "bern";
-    else
-        version = "unif";
+    int method = 1;
+    int n = 100;
+    double rate = 0.001;
+    double margin = 1;
+    int i;
+    if ((i = ArgPos((char *)"-size", argc, argv)) > 0) n = atoi(argv[i + 1]);
+    if ((i = ArgPos((char *)"-margin", argc, argv)) > 0) margin = atoi(argv[i + 1]);
+    if ((i = ArgPos((char *)"-method", argc, argv)) > 0) method = atoi(argv[i + 1]);
+    cout<<"size = "<<n<<endl;
+    cout<<"learing rate = "<<rate<<endl;
+    cout<<"margin = "<<margin<<endl;
+    if (method)
+        version = "bern";
+    else
+        version = "unif";
     cout<<"method = "<<version<<endl;
     prepare();
     train.run(n,rate,margin,method);