Increase robustness of is_split_into_words check: resolves ValueErr…

…or (#39)

* Increase robustness of "is_split_into_words" check

* Update the changelog

* Add a test case

CHANGELOG.md

@@ -20,6 +20,7 @@ Types of changes
 ### Fixed
 
 - No longer override `language` metadata from the dataset if the language was also set manually via `SpanMarkerModelCardData`.
+- No longer crash on `predict` with `ValueError: Failed to concatenate on axis=1 ...` if the first sentence in a list of sentences is just one word.
 
 ## [1.4.0]
 

span_marker/tokenizer.py

@@ -180,10 +180,14 @@ def __call__(
     ) -> Dict[str, List]:
         tokens = batch["tokens"]
         labels = batch.get("ner_tags", None)
-        # TODO: Increase robustness of this
         is_split_into_words = True
-        if isinstance(tokens, str) or (tokens and " " in tokens[0]):
+        if isinstance(tokens, str):
             is_split_into_words = False
+        elif tokens:
+            for token in tokens:
+                if " " in token:
+                    is_split_into_words = False
+                    break
 
         batch_encoding = self.tokenizer(
             tokens,

tests/test_modeling.py

@@ -195,6 +195,13 @@ def test_correct_predictions_with_document_level_context(
     )
 
 
+def test_predict_where_first_sentence_is_word(finetuned_conll_span_marker_model: SpanMarkerModel) -> None:
+    model = finetuned_conll_span_marker_model.try_cuda()
+    outputs = model.predict(["One", "Two Three Four Five"])
+    assert len(outputs) == 2
+    assert isinstance(outputs[0], list)
+
+
 def test_incorrect_predict_inputs(finetuned_conll_span_marker_model: SpanMarkerModel):
     model = finetuned_conll_span_marker_model.try_cuda()
     with pytest.raises(ValueError, match="could not recognize your input"):