Skip to content

Commit 0f0620b

Browse files
Fix SentencePiece version compatibility via hasattr check
1 parent 8066cf8 commit 0f0620b

1 file changed

Lines changed: 10 additions & 8 deletions

File tree

keras_hub/src/tokenizers/sentence_piece_tokenizer.py

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030

3131
VOCAB_FILENAME = "vocabulary.spm"
3232

33-
3433
@keras_hub_export("keras_hub.tokenizers.SentencePieceTokenizer")
3534
class SentencePieceTokenizer(tokenizer.Tokenizer):
3635
"""A SentencePiece tokenizer layer.
@@ -149,22 +148,25 @@ def _set_proto_tf(self, proto):
149148

150149
def _set_proto_spm(self, proto):
151150
out_type = str if is_string_dtype(self.compute_dtype) else int
152-
try:
153-
# Newer PyBind11 wrapper (>=0.2.2)
154-
self._sentence_piece_spm = spm.SentencePieceProcessor(
151+
152+
if hasattr(spm.SentencePieceProcessor(), "Init"):
153+
# Old SWIG wrapper (sentencepiece <=0.1.99)
154+
self._sentence_piece_spm = spm.SentencePieceProcessor()
155+
self._sentence_piece_spm.Init(
155156
model_proto=proto,
156157
out_type=out_type,
157158
add_bos=self.add_bos,
158159
add_eos=self.add_eos,
160+
alpha=1.0,
159161
)
160-
except TypeError:
161-
# Older SWIG wrapper (<0.2.2)
162-
self._sentence_piece_spm = spm.SentencePieceProcessor()
163-
self._sentence_piece_spm.Init(
162+
else:
163+
# New pybind11 wrapper (sentencepiece >=0.2.0)
164+
self._sentence_piece_spm = spm.SentencePieceProcessor(
164165
model_proto=proto,
165166
out_type=out_type,
166167
add_bos=self.add_bos,
167168
add_eos=self.add_eos,
169+
alpha=1.0,
168170
)
169171

170172
def set_proto(self, proto):

0 commit comments

Comments
 (0)