Fixes a bug which prevented the fragmenter from being able to process…

… tensors with a rank > 1.

PiperOrigin-RevId: 348670292

tensorflow_text/python/ops/state_based_sentence_breaker_op.py

@@ -87,7 +87,8 @@ def break_sentences_with_offsets(self, doc):
       end: A int64 `RaggedTensor` of shape [batch, (num_sentences)]
         where each entry is the exclusive ending byte offset of a sentence.
     """
-    if isinstance(doc, ragged_tensor.RaggedTensor):
+    if doc.shape.ndims > 1:
+      doc = ragged_tensor.RaggedTensor.from_tensor(doc)
       doc = doc.flat_values
 
     # Run sentence fragmenter op v2

tensorflow_text/python/ops/state_based_sentence_breaker_op_test.py

@@ -44,6 +44,13 @@ class SentenceFragmenterTestCasesV2(test.TestCase, parameterized.TestCase):
               b"Welcome to the U.S.", b"don't be surprised."
           ], [b"I.B.M.", b"yo"]],
       ),
+      dict(
+          test_description="Test when rank > 1.",
+          doc=[["Welcome to the U.S. don't be surprised."], ["I.B.M. yo"]],
+          expected_fragment_text=[[
+              b"Welcome to the U.S.", b"don't be surprised."
+          ], [b"I.B.M.", b"yo"]],
+      ),
       dict(
           test_description="Test semicolons",
           doc=["Welcome to the US; don't be surprised."],