Skip to content

Commit d10d1cd

Browse files
committed
Fixes a bug which prevented the fragmenter from being able to process tensors with a rank > 1.
PiperOrigin-RevId: 348670292
1 parent c1239e7 commit d10d1cd

File tree

2 files changed

+9
-1
lines changed

2 files changed

+9
-1
lines changed

tensorflow_text/python/ops/state_based_sentence_breaker_op.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,8 @@ def break_sentences_with_offsets(self, doc):
8787
end: A int64 `RaggedTensor` of shape [batch, (num_sentences)]
8888
where each entry is the exclusive ending byte offset of a sentence.
8989
"""
90-
if isinstance(doc, ragged_tensor.RaggedTensor):
90+
if doc.shape.ndims > 1:
91+
doc = ragged_tensor.RaggedTensor.from_tensor(doc)
9192
doc = doc.flat_values
9293

9394
# Run sentence fragmenter op v2

tensorflow_text/python/ops/state_based_sentence_breaker_op_test.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,13 @@ class SentenceFragmenterTestCasesV2(test.TestCase, parameterized.TestCase):
4444
b"Welcome to the U.S.", b"don't be surprised."
4545
], [b"I.B.M.", b"yo"]],
4646
),
47+
dict(
48+
test_description="Test when rank > 1.",
49+
doc=[["Welcome to the U.S. don't be surprised."], ["I.B.M. yo"]],
50+
expected_fragment_text=[[
51+
b"Welcome to the U.S.", b"don't be surprised."
52+
], [b"I.B.M.", b"yo"]],
53+
),
4754
dict(
4855
test_description="Test semicolons",
4956
doc=["Welcome to the US; don't be surprised."],

0 commit comments

Comments
 (0)