new: Add logic to handle single task case

Author: ivyleavedtoadflax

deep_reference_parser/deep_reference_parser.py

@@ -1035,11 +1035,16 @@ def predict(self, X, load_weights=False):
 
         # Compute the labels for each prediction for each task
 
+        # If running a single task model, wrap pred_index in a list so that it
+        # can use the same logic as multitask models.
+
+        if len(pred_index) == 1 :
+            pred_index = [pred_index]
+
         pred_label = []
         for i in range(len(ind2labelNew)):
             out = [[ind2labelNew[i][x] for x in a] for a in pred_index[i]]
             pred_label.append(out)
-
         # Flatten data
 
         # Remove the padded tokens. This is done by counting the number of

deep_reference_parser/parse.py

@@ -116,7 +116,7 @@ def parse(self, text, verbose=False):
 
         preds = self.drp.predict(tokens, load_weights=True)
 
-        flat_predictions = list(itertools.chain.from_iterable(preds))
+        flat_predictions = list(itertools.chain.from_iterable(preds))[0]
         flat_X = list(itertools.chain.from_iterable(tokens))
         rows = [i for i in zip(flat_X, flat_predictions)]
 

deep_reference_parser/split.py

@@ -124,7 +124,7 @@ def split(self, text, return_tokens=False, verbose=False):
 
         if return_tokens:
 
-            flat_predictions = list(itertools.chain.from_iterable(preds))
+            flat_predictions = list(itertools.chain.from_iterable(preds))[0]
             flat_X = list(itertools.chain.from_iterable(tokens))
             rows = [i for i in zip(flat_X, flat_predictions)]
 
@@ -145,7 +145,7 @@ def split(self, text, return_tokens=False, verbose=False):
 
             # Otherwise convert the tokens into references and return
 
-            refs = tokens_to_references(tokens, preds)
+            refs = tokens_to_references(tokens, preds[0])
 
             if verbose: