diff --git a/classifier/nn_ff.py b/classifier/nn_ff.py
index b8b8fc9..21c9024 100644
--- a/classifier/nn_ff.py
+++ b/classifier/nn_ff.py
@@ -1,9 +1,10 @@
-from utils.nn_reader import read_csv
+"""
+Simple feed-forward neural network in PyTorch for baseline results on Scicite data.
+Date: July 5th, 2020
+"""
 
 import torch
-
-
-
+from utils.nn_reader import read_csv_nn
 
 
 class Feedforward(torch.nn.Module):
@@ -26,63 +27,41 @@ class Feedforward(torch.nn.Module):
             return output
 
 
-"""
-from sklearn.datasets import make_blobs
-def blob_label(y, label, loc): # assign labels
-    target = numpy.copy(y)
-    for l in loc:
-        target[y == l] = label
-    return target
-X_train, y_train = make_blobs(n_samples=40, n_features=2, cluster_std=1.5, shuffle=True)
-X_train = torch.FloatTensor(X_train)
-y_train = torch.FloatTensor(blob_label(y_train, 0, [0]))
-y_train = torch.FloatTensor(blob_label(y_train, 1, [1,2,3]))
-x_test, y_test = make_blobs(n_samples=10, n_features=2, cluster_std=1.5, shuffle=True)
-x_test = torch.FloatTensor(x_test)
-y_test = torch.FloatTensor(blob_label(y_test, 0, [0]))
-y_test = torch.FloatTensor(blob_label(y_test, 1, [1,2,3]))
-"""
-
-
-X_train = torch.as_tensor(X_train)
-X_test = torch.as_tensor(X_test)
-y_train = torch.as_tensor(y_train)
-
-
-model = Feedforward(28, 9, 3)
-criterion = torch.nn.CrossEntropyLoss()
-optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)
-
 
+if __name__=='__main__':
+    X_train, y_train, X_test = read_csv_nn()
 
-model.eval()
-y_pred = model(X_train)
-y_pred = torch.Tensor([list(x).index(x.max()) for x in y_pred])
-y_pred = 
-before_train = criterion(y_train, y_pred)
-print('Test loss before training' , before_train.item())
+    X_train = torch.FloatTensor(X_train)
+    X_test = torch.FloatTensor(X_test)
+    y_train_ = torch.FloatTensor(y_train)
+    y_train = torch.max(torch.FloatTensor(y_train_),1)[1]
 
+    model = Feedforward(28, 9, 3)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr = 0.01)
 
-
-model.train()
-epoch = 20
-for epoch in range(epoch):
-    optimizer.zero_grad()
-    # Forward pass
+    model.eval()
     y_pred = model(X_train)
-    # Compute Loss
-    loss = criterion(y_pred.squeeze(), y_train)
-   
-    print('Epoch {}: train loss: {}'.format(epoch, loss.item()))
-    # Backward pass
-    loss.backward()
-    optimizer.step()
-
-
-model.eval()
-y_pred = model(X_test)
-after_train = criterion(y_pred.squeeze(), y_test) 
-print('Test loss after Training' , after_train.item())
+    before_train = criterion(y_pred, y_train)
+    print('Test loss before training' , before_train.item())
+
+    model.train()
+    epoch = 2000
+    for epoch in range(epoch):
+        optimizer.zero_grad()
+        # forward pass
+        y_pred = model(X_train)
+        loss = criterion(y_pred, y_train)
+       
+        print('Epoch {}: train loss: {}'.format(epoch, loss.item()))
+        # backward pass
+        loss.backward()
+        optimizer.step()
+
+    model.eval()
+    y_pred = model(X_train)
+    after_train = criterion(y_pred, y_train) 
+    print('Training loss after training' , after_train.item())
 
 
 
diff --git a/utils/nn_reader.py b/utils/nn_reader.py
index b68bc8e..647444d 100644
--- a/utils/nn_reader.py
+++ b/utils/nn_reader.py
@@ -1,43 +1,49 @@
 import numpy as np
-from iterables import chain
+from itertools import chain
 from utils.csv import read_csv_file
 
-train_file_path = 'data/tsv/train.tsv'
-test_file_path = 'data/tsv/test.tsv'
-train_raw = read_csv_file(train_file_path, '\t')
-
-features = [x.features for x in train_raw]
-features_unique = list(set(chain.from_iterable(features)))
-nobs = len(features)
-nfeats = len(features_unique)
-
-X_train = np.zeros((nobs, nfeats))
-
-for j in range(nfeats):
-    f = features_unique[j]
-    for i in range(nobs):
-        if f in features[i]:
-            X_train[i,j] = 1
-
-y_train_raw = np.array([x.true_label for x in train_raw])
-y_unique = sorted(list(set(y_train_raw)))
-y_dim = len(y_unique)
-y_train = np.zeros((nobs,y_dim))
-
-for j in range(y_dim):
-    y_train[:,j] = y_raw == y_unique[j]
-
-test_raw = read_csv_file(test_file_path, '\t')
-features = [x.features for x in test_raw]
-#features_unique = list(set(chain.from_iterable(features)))
-nobs = len(features)
-nfeats = len(features_unique)
-
-X_test = np.zeros((nobs, nfeats))
-for j in range(nfeats):
-    f = features_unique[j]
-    for i in range(nobs):
-        if f in features[i]:
-            X_test[i,j] = 1
+# TODO: clean up, transform into class, allow for command-line arguments
+
+def read_csv_nn(scicite_dir=None):
+
+	train_file_path = 'data/tsv/train.tsv'
+	test_file_path = 'data/tsv/test.tsv'
+	train_raw = read_csv_file(train_file_path, '\t')
+
+	features = [x.features for x in train_raw]
+	features_unique = list(set(chain.from_iterable(features)))
+	nobs = len(features)
+	nfeats = len(features_unique)
+
+	X_train = np.zeros((nobs, nfeats))
+
+	for j in range(nfeats):
+		f = features_unique[j]
+		for i in range(nobs):
+			if f in features[i]:
+				X_train[i,j] = 1
+
+	y_train_raw = np.array([x.true_label for x in train_raw])
+	y_unique = sorted(list(set(y_train_raw)))
+	y_dim = len(y_unique)
+	y_train = np.zeros((nobs,y_dim))
+
+	for j in range(y_dim):
+		y_train[:,j] = y_train_raw == y_unique[j]
+
+	test_raw = read_csv_file(test_file_path, '\t')
+	features = [x.features for x in test_raw]
+	#features_unique = list(set(chain.from_iterable(features)))
+	nobs = len(features)
+	nfeats = len(features_unique)
+
+	X_test = np.zeros((nobs, nfeats))
+	for j in range(nfeats):
+		f = features_unique[j]
+		for i in range(nobs):
+			if f in features[i]:
+				X_test[i,j] = 1
+
+	return X_train, y_train, X_test