Add out.tsv

main.py

@@ -23,27 +23,42 @@ def readInput(dir):
 def writeOutput(output, dir):
     with open(dir, 'w', newline='') as f:
         writer = csv.writer(f)
-        writer.writerows(output)
+        for row in output:
+            writer.writerow([row])
 
 
 if __name__ == '__main__':
-    train = pd.DataFrame(readInput('train/train.tsv.xz'), 
+    print('Reading input...')
+    train = pd.DataFrame(readInput('train/train.tsv.xz')[:10000], 
         columns=['Beginning', 'End', 'Title', 'Source', 'X'])
     train['Y'] = train.apply(lambda x: (float(x.Beginning) + float(x.End))/2, axis=1)
     train = train.drop(columns=['Beginning', 'End', 'Title', 'Source'])
 
     estimators = [('tfidf', TfidfVectorizer()), ('linearRegression', LinearRegression())]
+
+    print('Creating pipeline...')
     model = Pipeline(estimators)
+
+    print('Fitting model...')
     model.fit(train.X, train.Y)
 
+    print('dev-0...')
     # dev-0
     testX = readInput('dev-0/in.tsv')
-    writeOutput(model.predict(testX), 'dev-0/out.tsv')
+    predicted = model.predict(testX)
+    predicted = [str(x) for x in predicted.tolist()]
+    writeOutput(predicted, 'dev-0/out.tsv')
 
+    print('dev-1...')
     # dev-1
     testX = readInput('dev-1/in.tsv')
-    writeOutput(model.predict(testX), 'dev-1/out.tsv')
+    predicted = model.predict(testX)
+    predicted = [str(x) for x in predicted.tolist()]
+    writeOutput(predicted, 'dev-1/out.tsv')
 
+    print('test-A...')
     # test-A
     testX = readInput('test-A/in.tsv')
-    writeOutput(model.predict(testX), 'test-A/out.tsv')
+    predicted = model.predict(testX)
+    predicted = [str(x) for x in predicted.tolist()]
+    writeOutput(predicted, 'test-A/out.tsv')