Feat/chipper repetitions

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.7.24
+
+* Revised repetitions for Chipper
+
 ## 0.7.23
 
 * fix: added handling in `UnstructuredTableTransformerModel` for if `recognize` returns an empty

test_unstructured_inference/inference/test_layout.py

@@ -214,17 +214,14 @@ def points(self):
 
 class MockPageLayout(layout.PageLayout):
     def __init__(
-        self,
-        number=1,
-        image=None,
-        model=None,
-        detection_model=None,
+        self, number=1, image=None, model=None, detection_model=None, element_extraction_model=None
     ):
         self.image = image
         self.layout = layout
         self.model = model
         self.number = number
         self.detection_model = detection_model
+        self.element_extraction_model = element_extraction_model
 
 
 @pytest.mark.parametrize(

test_unstructured_inference/models/test_chippermodel.py

@@ -139,7 +139,7 @@ def test_no_repeat_ngram_logits():
 
     no_repeat_ngram_size = 2
 
-    logitsProcessor = chipper.NoRepeatNGramLogitsProcessor(ngram_size=2)
+    logitsProcessor = chipper.NoRepeatNGramLogitsProcessor(ngram_size=2, context_length=10)
     output = logitsProcessor(input_ids=input_ids, scores=logits)
 
     assert (
@@ -194,20 +194,64 @@ def test_ngram_repetiton_stopping_criteria():
     logits = torch.tensor([[0.1, -0.3, -0.5, 0, 1.0, -0.9]])
 
     stoppingCriteria = chipper.NGramRepetitonStoppingCriteria(
-        repetition_window=2, skip_tokens={0, 1, 2, 3, 4}
+        ngram_size=2, skip_tokens={0, 1, 2, 3, 4}
     )
 
     output = stoppingCriteria(input_ids=input_ids, scores=logits)
 
     assert output is False
 
     stoppingCriteria = chipper.NGramRepetitonStoppingCriteria(
-        repetition_window=2, skip_tokens={1, 2, 3, 4}
+        ngram_size=2, skip_tokens={1, 2, 3, 4}
     )
     output = stoppingCriteria(input_ids=input_ids, scores=logits)
     assert output is True
 
 
+def test_no_repeat_group_ngram_logits_processor():
+    input_ids = torch.tensor([[1, 2, 3, 4, 0, 1, 2, 3, 4]])
+    logits = torch.tensor([[0.1, -0.3, -0.5, 0, 1.0, -0.9]])
+
+    logitsProcessor = chipper.NoRepeatGroupNGramLogitsProcessor(ngram_size=3, token_group=[1, 2])
+
+    output = logitsProcessor(input_ids=input_ids, scores=logits)
+
+    assert (
+        int(
+            torch.sum(
+                output == torch.tensor([[0.1000, -0.3000, -0.5000, 0.0000, 1.0000, -0.9000]]),
+            ),
+        )
+        == 6
+    )
+
+    input_ids = torch.tensor([[1, 1, 2, 1, 2, 1, 2, 1, 2]])
+    logits = torch.tensor([[0.1, -0.3, -0.5, 0, 1.0, -0.9]])
+
+    output = logitsProcessor(input_ids=input_ids, scores=logits)
+
+    assert (
+        int(
+            torch.sum(
+                output
+                == torch.tensor([[0.1000, -float("inf"), -float("inf"), 0.0000, 1.0000, -0.9000]]),
+            ),
+        )
+        == 6
+    )
+
+
+def test_target_token_id_stopping_criterion():
+    input_ids = torch.tensor([1, 2, 3])
+    logits = torch.tensor([0.1, 0.2, 0.3])
+
+    stoppingCriterion = chipper.TargetTokenIdStoppingCriterion(1)
+
+    output = stoppingCriterion(input_ids=input_ids, scores=logits)
+
+    assert output is True
+
+
 @pytest.mark.parametrize(
     ("decoded_str", "expected_classes"),
     [
@@ -259,7 +303,7 @@ def test_predict_tokens_beam_indices():
     model = get_model("chipper")
     model.stopping_criteria = [
         chipper.NGramRepetitonStoppingCriteria(
-            repetition_window=1,
+            ngram_size=1,
             skip_tokens={},
         ),
     ]
@@ -294,9 +338,12 @@ def test_deduplicate_detected_elements():
     assert len(output) == 2
 
 
-def test_norepeatnGramlogitsprocessor_exception():
+def test_logitsprocessor_exception():
+    with pytest.raises(ValueError):
+        chipper.NoRepeatNGramLogitsProcessor(ngram_size="", context_length=10)
+
     with pytest.raises(ValueError):
-        chipper.NoRepeatNGramLogitsProcessor(ngram_size="")
+        chipper.NoRepeatGroupNGramLogitsProcessor(ngram_size="", token_group={})
 
 
 def test_run_chipper_v3():

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "0.7.23"  # pragma: no cover
+__version__ = "0.7.24"  # pragma: no cover

unstructured_inference/inference/layout.py

@@ -209,7 +209,9 @@ def get_elements_from_layout(
 
         # If the model is a chipper model, we don't want to order the
         # elements, as they are already ordered
-        order_elements = not isinstance(self.detection_model, UnstructuredChipperModel)
+        order_elements = not isinstance(
+            self.element_extraction_model, UnstructuredChipperModel
+        ) and not isinstance(self.detection_model, UnstructuredChipperModel)
         if order_elements:
             layout = order_layout(layout)