Update segmenter

data

@@ -1 +1 @@
-Subproject commit 33c57a3cafbdb46b4cc7db7f08695d63b52d6668
+Subproject commit dd266799aedd72e6381b368eacbe2767b6174aad

model.py

@@ -5,6 +5,8 @@ from torch.utils.data import Dataset, DataLoader
 import numpy as np
 from os import listdir
 from os.path import isfile, join
+import concurrent
+import itertools
 
 if __package__ == None or __package__ == "":
     from utils import tag_training_data, get_upenn_tags_dict, parse_tags    
@@ -79,20 +81,38 @@ class SegmentorDatasetDirectTag(Dataset):
 
 # The same dataset without one-hot embedding of the input.
 class SegmentorDatasetNonEmbed(Dataset):
+    @staticmethod
+    def read_file(f: str, document_root: str):
+        if f.endswith(".txt"):
+            fname = join(document_root, f)
+            print(f"Loaded datafile: {fname}")
+            reconstructed_tags = tag_training_data(fname)
+            input, tag = parse_tags(reconstructed_tags)
+            return [(
+                np.array(input),
+                np.array(tag)
+            )]
+        else:
+            return []
+        
     def __init__(self, document_root: str):
         self.datapoints = []
 
         files = listdir(document_root)
-        for f in files:
-            if f.endswith(".txt"):
-                fname = join(document_root, f)
-                print(f"Loaded datafile: {fname}")
-                reconstructed_tags = tag_training_data(fname)
-                input, tag = parse_tags(reconstructed_tags)
-                self.datapoints.append((
-                    np.array(input),
-                    np.array(tag)
-                ))
+        with concurrent.futures.ProcessPoolExecutor() as pool:
+            out = pool.map(SegmentorDatasetNonEmbed.read_file, files, itertools.repeat(document_root))
+        
+        self.datapoints = list(itertools.chain.from_iterable(out))
+        # for f in files:
+        #     if f.endswith(".txt"):
+        #         fname = join(document_root, f)
+        #         print(f"Loaded datafile: {fname}")
+        #         reconstructed_tags = tag_training_data(fname)
+        #         input, tag = parse_tags(reconstructed_tags)
+        #         self.datapoints.append((
+        #             np.array(input),
+        #             np.array(tag)
+        #         ))
         
     def __len__(self):
         return len(self.datapoints)

segmenter.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a8e6209584d0021684bb3a09ec1b717843f3086dfcc6411c57276f743f8e62fa
+oid sha256:005053e2036ac4a30364cdb81501140ef2ca238bee0f9a1a28fc5a4603d725f6
 size 10584544

train.py

@@ -26,6 +26,6 @@ if __name__ == "__main__":
     
     model.to(device)
 
-    train_bidirlstm_embedding_model(model, dataset, num_epochs=150, batch_size=2)
+    train_bidirlstm_embedding_model(model, dataset, num_epochs=100, batch_size=2)
 
     torch.save(model.state_dict(), "segmenter.ckpt")