refactoring tokens_to_references

Author: lizgzil

deep_reference_parser/tokens_to_references.py

@@ -10,21 +10,12 @@
 from .deep_reference_parser import logger
 
 
-def tokens_to_references(tokens, labels):
-    """
-    Given a corresponding list of tokens and a list of labels: split the tokens
-    and return a list of references.
-
-    Args:
-        tokens(list): A list of tokens.
-        labels(list): A corresponding list of labels.
-
-    """
+def get_reference_spans(tokens, spans):
 
     # Flatten the lists of tokens and predictions into a single list.
 
     flat_tokens = list(itertools.chain.from_iterable(tokens))
-    flat_predictions = list(itertools.chain.from_iterable(labels))
+    flat_predictions = list(itertools.chain.from_iterable(spans))
 
     # Find all b-r and e-r tokens.
 
@@ -37,30 +28,43 @@ def tokens_to_references(tokens, labels):
     logger.debug("Found %s b-r tokens", len(ref_starts))
     logger.debug("Found %s e-r tokens", len(ref_ends))
 
-    references = []
-
     n_refs = len(ref_starts)
 
     # Split on each b-r.
-    # TODO: It may be worth including some simple post processing steps here
-    # to pick up false positives, for instance cutting short a reference
-    # after n tokens.
 
+    token_starts = []
+    token_ends = []
     for i in range(0, n_refs):
-        token_start = ref_starts[i]
+        token_starts.append(ref_starts[i])
         if i + 1 < n_refs:
-
-            token_end = ref_starts[i + 1] - 1
+            token_ends.append(ref_starts[i + 1] - 1)
         else:
-            token_end = len(flat_tokens)
+            token_ends.append(len(flat_tokens))
 
+    return token_starts, token_ends, flat_tokens
+
+
+def tokens_to_references(tokens, labels):
+    """
+    Given a corresponding list of tokens and a list of labels: split the tokens
+    and return a list of references.
+
+    Args:
+        tokens(list): A list of tokens.
+        labels(list): A corresponding list of labels.
+
+    """
+
+    token_starts, token_ends, flat_tokens = get_reference_spans(tokens, labels)
+
+    references = []
+    for token_start, token_end in zip(token_starts, token_ends):
         ref = flat_tokens[token_start : token_end + 1]
         flat_ref = " ".join(ref)
         references.append(flat_ref)
 
     return references
 
-
 def tokens_to_reference_lists(tokens, spans, components):
     """
     Given a corresponding list of tokens, a list of
@@ -75,37 +79,12 @@ def tokens_to_reference_lists(tokens, spans, components):
 
     """
 
-    # Flatten the lists of tokens and predictions into a single list.
-
-    flat_tokens = list(itertools.chain.from_iterable(tokens))
-    flat_spans = list(itertools.chain.from_iterable(spans))
+    token_starts, token_ends, flat_tokens = get_reference_spans(tokens, spans)
     flat_components = list(itertools.chain.from_iterable(components))
 
-    # Find all b-r and e-r tokens.
-
-    ref_starts = [
-        index for index, label in enumerate(flat_spans) if label == "b-r"
-    ]
-
-    ref_ends = [index for index, label in enumerate(flat_spans) if label == "e-r"]
-
-    logger.debug("Found %s b-r tokens", len(ref_starts))
-    logger.debug("Found %s e-r tokens", len(ref_ends))
-
     references_components = []
+    for token_start, token_end in zip(token_starts, token_ends):
 
-    n_refs = len(ref_starts)
-
-    # Split on each b-r.
-
-    for i in range(0, n_refs):
-        token_start = ref_starts[i]
-        if i + 1 < n_refs:
-
-            token_end = ref_starts[i + 1] - 1
-        else:
-            token_end = len(flat_tokens)
-        
         ref_tokens = flat_tokens[token_start : token_end + 1]
         ref_components = flat_components[token_start : token_end + 1]
         flat_ref = " ".join(ref_tokens)